diff --git a/.github/vllm-main-verified.commit b/.github/vllm-main-verified.commit index aef2c04f389a..f0b59d1e8aa2 100644 --- a/.github/vllm-main-verified.commit +++ b/.github/vllm-main-verified.commit @@ -1 +1 @@ -24420b7a341af37a9bc3186e6fc6062512e0aa58 +3b95523ed904cf02a26b872f4a900ef17e6be8ad diff --git a/docs/source/user_guide/feature_guide/dyntra_lb.md b/docs/source/user_guide/feature_guide/dyntra_lb.md index ebc2e368b0d7..3902b7ea5739 100644 --- a/docs/source/user_guide/feature_guide/dyntra_lb.md +++ b/docs/source/user_guide/feature_guide/dyntra_lb.md @@ -40,9 +40,15 @@ DyntraLB cannot be combined with: - `scheduler_config.enable_balance_scheduling` - `scheduler_config.profiling_chunk_config.enabled` +- vLLM main's `max_async_kv_loads_per_step` The service rejects these unsupported combinations during startup. +The async KV load limit must also remain unset with recompute, profiling-chunk, +or balance scheduling. These Ascend schedulers do not yet consume the separate +async waiting queue used by this upstream option. Schedulers that inherit +upstream scheduling can use the limit. + ## Configuration Add `dyntra_lb_config` under `scheduler_config` in the Decode node's diff --git a/examples/offline_DSD_k_tuner/generate_config.py b/examples/offline_DSD_k_tuner/generate_config.py index d719f30c77a2..182bafe46b42 100644 --- a/examples/offline_DSD_k_tuner/generate_config.py +++ b/examples/offline_DSD_k_tuner/generate_config.py @@ -17,6 +17,8 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.metrics.reader import Counter, Vector +from vllm_ascend.utils import vllm_version_is + DEFAULT_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 6, 8, 16, 24, 32, 64, 96, 128, 256, 384, 512) @@ -234,6 +236,9 @@ def run_profiling_sweep(args): def get_acceptance_rate_per_pos(args): """Get acceptance rate per position.""" + if args.method in ("eagle", "eagle3") and args.disable_padded_drafter_batch and not vllm_version_is("0.33.0"): + raise ValueError("--disable-padded-drafter-batch is supported only with vLLM v0.33.0.") + tokenizer = AutoTokenizer.from_pretrained(args.model_dir) prompts = get_samples(args, tokenizer) llm_prompts: list[Any] @@ -260,7 +265,9 @@ def get_acceptance_rate_per_pos(args): "method": args.method, "model": eagle_dir, "num_speculative_tokens": args.num_spec_tokens, - "disable_padded_drafter_batch": args.disable_padded_drafter_batch, + **( + {"disable_padded_drafter_batch": args.disable_padded_drafter_batch} if vllm_version_is("0.33.0") else {} + ), "parallel_drafting": args.parallel_drafting, } elif args.method == "ngram": @@ -387,7 +394,11 @@ def main(): parser.add_argument("--enforce-eager", action="store_true") parser.add_argument("--enable-chunked-prefill", action="store_true") parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) - parser.add_argument("--disable-padded-drafter-batch", action="store_true") + parser.add_argument( + "--disable-padded-drafter-batch", + action="store_true", + help="Disable drafter padding for Eagle/Eagle3 on vLLM v0.33.0.", + ) parser.add_argument("--result-dir", type=str, default="./log/dynamic_sd") parser.add_argument("--prompt-lookup-max", type=int, default=5) parser.add_argument("--prompt-lookup-min", type=int, default=2) diff --git a/release_notes/fragments/21469.yaml b/release_notes/fragments/21469.yaml new file mode 100644 index 000000000000..5ab220bdec2d --- /dev/null +++ b/release_notes/fragments/21469.yaml @@ -0,0 +1,6 @@ +pr: 21469 +entries: + - category: dependency + text: Updated the verified vLLM main revision from the November 6 pin to the November 12 pin while retaining v0.33.0 compatibility. + - category: bugfix + text: Adapted Ascend to the intervening compiler, scheduler, model runner, sampling, MoE scaling, pipeline auxiliary-state, and sparse-attention API changes. diff --git a/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_rejection_sample_v2.py b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_rejection_sample_v2.py index 05e1b2de6694..592ddb0dc4b3 100644 --- a/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_rejection_sample_v2.py +++ b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_rejection_sample_v2.py @@ -3,17 +3,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc +from importlib import import_module +from typing import TYPE_CHECKING import numpy as np import pytest import torch from vllm.triton_utils import tl, triton -from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates from vllm_ascend.device.device_config import get_ascend_device_type from vllm_ascend.device.hardware import AscendDeviceType +from vllm_ascend.utils import vllm_version_is from vllm_ascend.worker.v2.spec_decode.rejection_sampler_utils import rejection_sample +if TYPE_CHECKING or not vllm_version_is("0.33.0"): + from vllm.v1.worker.gpu.spec_decode.rejection_sampler import unconditional_to_conditional_rates +else: + unconditional_to_conditional_rates = import_module("vllm.v1.spec_decode.utils").unconditional_to_conditional_rates + VOCAB_SIZE = 4096 pytest.importorskip("triton") diff --git a/tests/e2e/pull_request/four_card/context_parallel/test_accuracy_v2.py b/tests/e2e/pull_request/four_card/context_parallel/test_accuracy_v2.py index 4ee67b14932e..fe930e2864dd 100644 --- a/tests/e2e/pull_request/four_card/context_parallel/test_accuracy_v2.py +++ b/tests/e2e/pull_request/four_card/context_parallel/test_accuracy_v2.py @@ -31,6 +31,7 @@ from tests.e2e.conftest import DPVllmRunner, VllmRunner, wait_until_npu_memory_free from tests.e2e.pull_request.one_card.model_runner_v2.utils import calculate_acceptance_per_pos from tests.e2e.pull_request.utils import run_pd_disaggregation +from vllm_ascend.utils import vllm_version_is MAX_NUM_SEQS = 4 FULL_DECODE_GRAPH = { @@ -583,7 +584,7 @@ def test_dsv3_2_sfa_pcp_dcp_model_runner_v2_graph_accuracy() -> None: "model": EAGLE3_PCP_DRAFT_MODEL, "num_speculative_tokens": 3, "draft_tensor_parallel_size": 2, - "disable_padded_drafter_batch": False, + **({"disable_padded_drafter_batch": False} if vllm_version_is("0.33.0") else {}), }, }, ) diff --git a/tests/e2e/pull_request/four_card/test_sparse_hccl_checkpoint.py b/tests/e2e/pull_request/four_card/test_sparse_hccl_checkpoint.py index df8e0d3e2c33..c6db8afc5317 100644 --- a/tests/e2e/pull_request/four_card/test_sparse_hccl_checkpoint.py +++ b/tests/e2e/pull_request/four_card/test_sparse_hccl_checkpoint.py @@ -302,6 +302,11 @@ def test_sparse_checkpoint_matches_native_dense(tmp_path, tp_size, dp_size): indent=2, ) ) + except BaseException: + # Runtime RPC failures otherwise expose only HTTP 500; retain the + # server's original traceback before the temporary directory exits. + print((tmp_path / "server.log").read_text(errors="replace"), flush=True) + raise finally: if engine is not None: engine.shutdown() diff --git a/tests/e2e/pull_request/one_card/spec_decode/test_eagle.py b/tests/e2e/pull_request/one_card/spec_decode/test_eagle.py index 8f2e7408b216..c56dccdcb937 100644 --- a/tests/e2e/pull_request/one_card/spec_decode/test_eagle.py +++ b/tests/e2e/pull_request/one_card/spec_decode/test_eagle.py @@ -9,6 +9,7 @@ from tests.e2e.conftest import VllmRunner from tests.e2e.pull_request.one_card.spec_decode.utils import BASELINES, MODELS, calculate_acceptance_per_pos +from vllm_ascend.utils import vllm_version_is MAX_NUM_SEQS = 256 @@ -67,7 +68,7 @@ def test_qwen_eagle3_acceptance( "method": method, "num_speculative_tokens": num_speculative_tokens, "draft_tensor_parallel_size": draft_tensor_parallel_size, - "disable_padded_drafter_batch": disable_padded_drafter_batch, + **({"disable_padded_drafter_batch": disable_padded_drafter_batch} if vllm_version_is("0.33.0") else {}), "model": spec_model_name, } if dynamic_num_speculative_tokens is not None: diff --git a/tests/e2e/pull_request/one_card/spec_decode/test_mtp_eagle_correctness.py b/tests/e2e/pull_request/one_card/spec_decode/test_mtp_eagle_correctness.py index 430cb53ac861..13c72d67280e 100644 --- a/tests/e2e/pull_request/one_card/spec_decode/test_mtp_eagle_correctness.py +++ b/tests/e2e/pull_request/one_card/spec_decode/test_mtp_eagle_correctness.py @@ -27,6 +27,7 @@ from vllm.config import CompilationConfig from tests.e2e.conftest import VllmRunner, cleanup_dist_env_and_memory +from vllm_ascend.utils import vllm_version_is os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -65,7 +66,7 @@ def test_deepseek_mtp( speculative_config = { "method": "mtp", "num_speculative_tokens": num_speculative_tokens, - "disable_padded_drafter_batch": disable_padded_drafter_batch, + **({"disable_padded_drafter_batch": disable_padded_drafter_batch} if vllm_version_is("0.33.0") else {}), } if dynamic_num_speculative_tokens is not None: speculative_config["num_speculative_tokens_per_batch_size"] = [ diff --git a/tests/e2e/pull_request/one_card/test_categorical_sample_bias.py b/tests/e2e/pull_request/one_card/test_categorical_sample_bias.py new file mode 100644 index 000000000000..c67d8762739b --- /dev/null +++ b/tests/e2e/pull_request/one_card/test_categorical_sample_bias.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch_npu # noqa: F401 + +from vllm_ascend.ops.triton.sample.categorical_sample import categorical_sample +from vllm_ascend.ops.triton.triton_utils import init_device_properties_triton + +pytestmark = pytest.mark.e2e_coverage( + arch="", feature="dspark", parallel="", deploy="", hardware="A2", quantization="", graph_mode="eager" +) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_markov_bias_sampling_and_draft_cache_on_npu(dtype, with_bias): + """Exercise the real kernel, row mapping and fine/coarse tail boundaries.""" + init_device_properties_triton() + logits_cpu = torch.zeros(2, 8193, dtype=dtype) + logits_cpu[:, 0] = 1 + bias_cpu = torch.zeros_like(logits_cpu) + bias_cpu[0, 8192] = 4 + bias_cpu[1, 1024] = 4 + expected = (logits_cpu.float() + bias_cpu.float()).to(dtype) if with_bias else logits_cpu + logits = logits_cpu.to("npu") + bias = bias_cpu.to("npu") if with_bias else None + cache = torch.full((3, 3, 8193), -33, dtype=dtype, device="npu") + sampled = categorical_sample( + logits, + torch.tensor([1, 0], dtype=torch.int32, device="npu"), + torch.zeros(3, device="npu"), + torch.tensor([3, 7, 11], dtype=torch.int64, device="npu"), + torch.tensor([8, 12], dtype=torch.int64, device="npu"), + apply_temperature=True, + is_drafting=True, + logits_cache=cache, + logits_cache_col=torch.tensor([1, 2], dtype=torch.int32, device="npu"), + bias=bias, + ) + torch.npu.synchronize() + torch.testing.assert_close(sampled.cpu(), expected.argmax(dim=-1), rtol=0, atol=0) + torch.testing.assert_close(cache[1, 1].cpu(), expected[0], rtol=0, atol=0) + torch.testing.assert_close(cache[0, 2].cpu(), expected[1], rtol=0, atol=0) + torch.testing.assert_close(cache[2].cpu(), torch.full((3, 8193), -33, dtype=dtype), rtol=0, atol=0) + torch.testing.assert_close(logits.cpu(), logits_cpu, rtol=0, atol=0) diff --git a/tests/e2e/pull_request/one_card/test_qwen3_8b_w8a8.py b/tests/e2e/pull_request/one_card/test_qwen3_8b_w8a8.py index ce5b27dc1d4e..4f8d87beeadd 100644 --- a/tests/e2e/pull_request/one_card/test_qwen3_8b_w8a8.py +++ b/tests/e2e/pull_request/one_card/test_qwen3_8b_w8a8.py @@ -21,6 +21,7 @@ from tests.e2e.conftest import VllmRunner, cleanup_dist_env_and_memory, wait_until_npu_memory_free from tests.e2e.pull_request.utils import PROMPTS_SHORT +from vllm_ascend.utils import vllm_version_is @wait_until_npu_memory_free() @@ -44,7 +45,7 @@ def test_dense_w8a8_eagle3_full_graph(): async_scheduling=False, quantization="ascend", speculative_config={ - "disable_padded_drafter_batch": False, + **({"disable_padded_drafter_batch": False} if vllm_version_is("0.33.0") else {}), "method": "eagle3", "model": "RedHatAI/Qwen3-8B-speculator.eagle3", "num_speculative_tokens": 2, diff --git a/tests/e2e/pull_request/two_card/spec_decode/test_spec_decode.py b/tests/e2e/pull_request/two_card/spec_decode/test_spec_decode.py index f8f95a77c665..27718d58d992 100644 --- a/tests/e2e/pull_request/two_card/spec_decode/test_spec_decode.py +++ b/tests/e2e/pull_request/two_card/spec_decode/test_spec_decode.py @@ -32,6 +32,7 @@ from vllm.v1.metrics.reader import Counter, Vector from tests.e2e.conftest import VllmRunner +from vllm_ascend.utils import vllm_version_is os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -119,7 +120,9 @@ def test_eagle3_sp_acceptance( "enforce_eager": True, "method": method, "num_speculative_tokens": num_speculative_tokens, - "disable_padded_drafter_batch": disable_padded_drafter_batch, + # Main removed this MRv1 option; both retained parameter identities + # exercise native MRv2 padding, not legacy unpadded-drafter coverage. + **({"disable_padded_drafter_batch": disable_padded_drafter_batch} if vllm_version_is("0.33.0") else {}), "model": spec_model_name, } diff --git a/tests/ut/_310p/ops/test_mrope_state_310.py b/tests/ut/_310p/ops/test_mrope_state_310.py new file mode 100644 index 000000000000..06cdd3583ed7 --- /dev/null +++ b/tests/ut/_310p/ops/test_mrope_state_310.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU validation of persistent 310P MRoPE host-buffer updates.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_ascend._310p.worker.v2.rope import Ascend310PRopeState + + +@pytest.mark.parametrize("num_dims", [3, 4]) +def test_mrope_mixed_prefill_decode_preserves_padding_and_dummy_column(num_dims): + state = object.__new__(Ascend310PRopeState) + state.num_dims = num_dims + state.positions_cpu = torch.full((num_dims, 9), -71, dtype=torch.int64) + state.positions_np = state.positions_cpu.numpy() + state.positions = torch.full_like(state.positions_cpu, -90) + state.prefill_delta = SimpleNamespace(np=np.array([5, 0, -2])) + state.prefill_positions = SimpleNamespace(cpu=torch.arange(3 * num_dims * 12).view(3 * num_dims, 12)) + state.prepare_positions_cpu( + np.array([2, 0, 1]), np.array([0, 2, 5, 6]), np.array([0, 10, 0]), np.array([7, 2, 4]), 8 + ) + expected = torch.zeros((num_dims, 9), dtype=torch.int64) + expected[:, :2] = torch.tensor([2, 3]) + expected[:, 2:5] = torch.tensor([12, 13, 14]) + expected[:, 5] = state.prefill_positions.cpu[num_dims : 2 * num_dims, 2] + expected[:, 8] = -71 + torch.testing.assert_close(state.positions_cpu, expected, rtol=0, atol=0) + torch.testing.assert_close(state.positions, expected, rtol=0, atol=0) diff --git a/tests/ut/attention/test_flash_attn.py b/tests/ut/attention/test_flash_attn.py index dd1648409734..b8458c47cc41 100644 --- a/tests/ut/attention/test_flash_attn.py +++ b/tests/ut/attention/test_flash_attn.py @@ -15,6 +15,7 @@ from vllm_ascend.attention.attention_v1 import AscendAttentionBackend from vllm_ascend.attention.fa_v1 import AscendFlashAttentionBackend, AscendFlashAttentionImpl from vllm_ascend.models.minimax_m3 import msa_m3 +from vllm_ascend.utils import vllm_version_is from vllm_ascend.worker.v2 import attn_utils @@ -389,7 +390,12 @@ def forbidden_host_materialization(*args, **kwargs): with monkeypatch.context() as patches: for operation in ("tolist", "item", "cpu"): patches.setattr(torch.Tensor, operation, forbidden_host_materialization) - result = builder.build_for_drafting(common, draft_index=0) + if vllm_version_is("0.33.0"): + result = builder.build_for_drafting(common, draft_index=0) + else: + # MRv2 builds each draft step directly; the V1 drafting helper was + # removed upstream. Retain the same device-length assertions. + result = builder.build(common_prefix_len=0, common_attn_metadata=common, fast_build=True) assert result.seq_lens_gpu.data_ptr() == common.seq_lens.data_ptr() assert result.query_start_loc_gpu.data_ptr() == common.query_start_loc.data_ptr() assert not result.causal diff --git a/tests/ut/attention/test_mla_v1.py b/tests/ut/attention/test_mla_v1.py index 5bef5d314ddd..958fa00e33bd 100644 --- a/tests/ut/attention/test_mla_v1.py +++ b/tests/ut/attention/test_mla_v1.py @@ -6,7 +6,7 @@ import torch from vllm.compilation import backends from vllm.compilation.backends import set_model_tag -from vllm.config import CacheConfig, SchedulerConfig, VllmConfig +from vllm.config import CacheConfig, SchedulerConfig, SpeculativeConfig, VllmConfig from vllm.distributed.parallel_state import GroupCoordinator from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod @@ -34,6 +34,19 @@ from vllm_ascend.utils import ACL_FORMAT_FRACTAL_ND, ACL_FORMAT_FRACTAL_NZ, vllm_version_is +def test_mla_block_table_uses_actual_native_padding_default(): + speculative = object.__new__(SpeculativeConfig) + if vllm_version_is("0.33.0"): + speculative.disable_padded_drafter_batch = False + builder = object.__new__(AscendMLAMetadataBuilder) + builder.speculative_config = speculative + builder.graph_pad_size = 8 + builder.num_decodes = 2 + metadata = SimpleNamespace(num_reqs=4) + assert builder.get_block_table_size(metadata, 0) == 4 + assert builder.get_block_table_size(metadata, 1) == 2 + + @pytest.mark.parametrize("num_tokens", [1, 3]) @pytest.mark.parametrize( "num_heads,kv_lora_rank", @@ -1116,7 +1129,10 @@ def test_ascend_mla_metadata_builder_build_full_graph(self, mock_get_cos_and_sin block_table = torch.Tensor([[1, 0], [2, 0], [3, 0], [4, 0]]).int() common_metadata.block_table_tensor = block_table mock_get_cos_and_sin_mla.return_value = (torch.tensor([6, 6]), torch.Tensor([6, 6])) - metadata = builder.build(0, common_metadata) + # Preserve this original MRv1 unpadded-layout contract on both lanes; + # main's actual native default is exercised separately above. + with patch("vllm_ascend.attention.mla_v1.vllm_version_is", return_value=True): + metadata = builder.build(0, common_metadata) self.assertFalse(metadata.causal) self.assertEqual(metadata.decode.actual_seq_lengths_q, [1, 2, 4, 5, 6, 6, 7, 8]) diff --git a/tests/ut/core/test_dyntra_lb_scheduler.py b/tests/ut/core/test_dyntra_lb_scheduler.py index 955da04fc262..942bd716c040 100644 --- a/tests/ut/core/test_dyntra_lb_scheduler.py +++ b/tests/ut/core/test_dyntra_lb_scheduler.py @@ -1,4 +1,5 @@ import time +from copy import copy from types import SimpleNamespace from typing import Any, TypeVar from unittest.mock import patch @@ -32,6 +33,7 @@ get_dyntra_lb_request_block_num, print_scheduler_summary, ) +from vllm_ascend.utils import vllm_version_is SchedulerT = TypeVar("SchedulerT", bound=Scheduler) @@ -678,6 +680,11 @@ def build_metadata(connector, output): scheduler_output = scheduler.schedule() + if not vllm_version_is("0.33.0"): + # Main removed encoder metadata from its wire payload. Observe the + # manager's real extension payload separately for the original checks. + scheduler_output.ec_manager_metadata = scheduler.encoder_cache_manager.get_manager_metadata() + assert len(block_states) == 1 assert block_states[0].boundary_state_offloads is boundary_state_offloads assert block_states[0].req_ids == set() @@ -1023,8 +1030,10 @@ def test_dyntra_lb_async_v2_sets_pp_decode_cadence(): scheduler_output = scheduler.schedule() assert request.request_id in scheduler_output.num_scheduled_tokens - assert scheduler.pp_size == 2 - assert request.next_decode_eligible_step == scheduler.current_step + scheduler.pp_size + assert scheduler.parallel_config.pipeline_parallel_size == 2 + assert ( + request.next_decode_eligible_step == scheduler.current_step + scheduler.parallel_config.pipeline_parallel_size + ) def test_dyntra_lb_sync_skips_until_previous_output_arrives(): @@ -1097,6 +1106,20 @@ def test_dyntra_lb_paused_request_resumes_as_cached(): resumed_output = scheduler.schedule() + if not vllm_version_is("0.33.0"): + # Native MRv2 transports readmitted requests as new payloads. Verify + # those real payloads before exposing the old classification view used + # by the original pause/resume assertions below. + assert [data.req_id for data in resumed_output.scheduled_new_reqs] == [request.request_id] + assert resumed_output.scheduled_new_reqs[0].block_ids == block_ids + assert resumed_output.scheduled_cached_reqs.req_ids == [] + resumed_output = copy(resumed_output) + resumed_output.scheduled_cached_reqs = copy(resumed_output.scheduled_cached_reqs) + resumed_output.scheduled_cached_reqs.resumed_req_ids = { + data.req_id for data in resumed_output.scheduled_new_reqs + } + resumed_output.scheduled_new_reqs = [] + assert request.status == RequestStatus.RUNNING assert request.request_id not in scheduler._lb_paused_req_ids assert request.request_id in resumed_output.scheduled_cached_reqs.resumed_req_ids diff --git a/tests/ut/core/test_profiling_chunk.py b/tests/ut/core/test_profiling_chunk.py index b77e033cb90b..1b70e615145e 100644 --- a/tests/ut/core/test_profiling_chunk.py +++ b/tests/ut/core/test_profiling_chunk.py @@ -37,6 +37,7 @@ from vllm_ascend.core.short_request_first_scheduler import ( ShortRequestFirstRequestQueue, ) +from vllm_ascend.utils import vllm_version_is MODEL = "Qwen/Qwen3-0.6B" BLOCK_SIZE = 16 @@ -67,7 +68,7 @@ def make_output(scheduler): sampled_token_ids = [[1000]] * len(scheduler.running) return ModelRunnerOutput( req_ids=req_ids, - req_id_to_index=req_id_to_index, + **({"req_id_to_index": req_id_to_index} if vllm_version_is("0.33.0") else {}), sampled_token_ids=sampled_token_ids, logprobs=None, prompt_logprobs_dict={}, @@ -982,10 +983,10 @@ def test_v2_pp_sets_next_decode_eligible_step(self): first_output = scheduler.schedule() self.assertIn(request.request_id, first_output.num_scheduled_tokens) - self.assertEqual(scheduler.pp_size, scheduler.parallel_config.pipeline_parallel_size) + self.assertEqual(scheduler.parallel_config.pipeline_parallel_size, 2) self.assertEqual( request.next_decode_eligible_step, - scheduler.current_step + scheduler.pp_size, + scheduler.current_step + scheduler.parallel_config.pipeline_parallel_size, ) # The next step is inside the PP cadence window: the request is skipped. diff --git a/tests/ut/core/test_recompute_scheduler.py b/tests/ut/core/test_recompute_scheduler.py index c9d3344e4353..b6f3ead02310 100644 --- a/tests/ut/core/test_recompute_scheduler.py +++ b/tests/ut/core/test_recompute_scheduler.py @@ -24,6 +24,7 @@ _normalize_encoder_schedule_result, ) from vllm_ascend.ec_manager.score_ec_manager import ScoreEncoderCacheManager +from vllm_ascend.utils import vllm_version_is @pytest.fixture(autouse=True) @@ -409,6 +410,9 @@ def test_waiting_multimodal_duplicate_cache_tracking(): manager = scheduler.encoder_cache_manager assert manager.request_cached_ids[request.request_id] == {0, 1} assert manager.cpu_num_free_slots == 6 + if not vllm_version_is("0.33.0"): + # Diagnostic manager payload, not a field transported by native MRv2. + output.ec_manager_metadata = manager.get_manager_metadata() assert output.ec_manager_metadata.cpu_get_encoder_mm_hashes == [] manager._check_invariant() scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) diff --git a/tests/ut/core/test_scheduler_connector_block_state.py b/tests/ut/core/test_scheduler_connector_block_state.py index d51fdc420271..a09efce11af3 100644 --- a/tests/ut/core/test_scheduler_connector_block_state.py +++ b/tests/ut/core/test_scheduler_connector_block_state.py @@ -43,7 +43,7 @@ def test_boundary_state_is_drained_consumed_and_not_dispatched(monkeypatch, sche scheduler.ec_connector = None scheduler.requests = {"cached": object(), "boundary": object()} cached_data = SimpleNamespace(req_ids=["cached", "unchanged"], new_block_ids=[([9],), None]) - monkeypatch.setattr(scheduler, "_make_cached_request_data", lambda *args: cached_data) + monkeypatch.setattr(scheduler, "_make_cached_request_data", lambda *args, **kwargs: cached_data) offers = {"boundary": [(0, 42, 128)], "finished": [(0, 43, 128)]} drain = Mock(side_effect=[offers, {}]) get_blocks = Mock(side_effect=lambda req_id: {"cached": ([1, 9],), "boundary": ([42],)}[req_id]) diff --git a/tests/ut/core/test_score_encoder_cache_scheduler.py b/tests/ut/core/test_score_encoder_cache_scheduler.py index a1ff0b98c3a4..cb45f872e2dd 100644 --- a/tests/ut/core/test_score_encoder_cache_scheduler.py +++ b/tests/ut/core/test_score_encoder_cache_scheduler.py @@ -14,6 +14,7 @@ from tests.ut.kv_offload.utils import create_model_runner_output, create_scheduler, create_vllm_config from vllm_ascend.ec_manager.score_ec_manager import ScoreEncoderCacheManager +from vllm_ascend.utils import vllm_version_is def _make_scheduler(npu_slots: int = 8) -> Scheduler: @@ -87,7 +88,10 @@ def _capacity(manager: ScoreEncoderCacheManager) -> tuple[int, int, int, int]: def _assert_error_and_finish(scheduler: Scheduler, request: Request, scheduler_output) -> None: assert scheduler_output.total_num_scheduled_tokens == 0 assert not scheduler_output.scheduled_encoder_inputs - outputs = scheduler.update_from_output(scheduler_output, ModelRunnerOutput(req_ids=[], req_id_to_index={})) + outputs = scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput(req_ids=[], **({"req_id_to_index": {}} if vllm_version_is("0.33.0") else {})), + ) assert request.status == RequestStatus.FINISHED_ERROR assert request.request_id not in scheduler.requests assert request.request_id not in scheduler.encoder_cache_manager.request_cached_ids diff --git a/tests/ut/distributed/ascend_store/test_layerwise_keys.py b/tests/ut/distributed/ascend_store/test_layerwise_keys.py index d2edb3ab0a9a..aa742105c061 100644 --- a/tests/ut/distributed/ascend_store/test_layerwise_keys.py +++ b/tests/ut/distributed/ascend_store/test_layerwise_keys.py @@ -147,8 +147,14 @@ def make_config(pp_size=1): get_layers_start_end_indices=lambda p: ((0, 3), (3, 7))[p.rank // 2], get_total_num_hidden_layers=lambda: 7, compute_hash=lambda: "model-config", + compile_factors=lambda: {"identity": "model-config"}, + ), + cache_config=SimpleNamespace( + block_size=16, + cache_dtype="auto", + compute_hash=lambda: "cache-config", + compile_factors=lambda: {"identity": "cache-config"}, ), - cache_config=SimpleNamespace(block_size=16, cache_dtype="auto", compute_hash=lambda: "cache-config"), speculative_config=None, ) diff --git a/tests/ut/distributed/ascend_store/test_mooncake_pipeline.py b/tests/ut/distributed/ascend_store/test_mooncake_pipeline.py index 21330068fb19..46c4b14a9f52 100644 --- a/tests/ut/distributed/ascend_store/test_mooncake_pipeline.py +++ b/tests/ut/distributed/ascend_store/test_mooncake_pipeline.py @@ -298,6 +298,9 @@ def compute_hash(self): # Upstream CacheConfig.compute_hash includes block_size. return f"{self.block_size}:{self.cache_dtype}" + def compile_factors(self): + return {"block_size": self.block_size, "cache_dtype": self.cache_dtype} + worker, _ = self.make_stage(0, MemoryRangeStore()) config = worker.vllm_config config.cache_config = CacheConfig(32) diff --git a/tests/ut/distributed/ascend_store/test_pool_worker.py b/tests/ut/distributed/ascend_store/test_pool_worker.py index 12d11c90bfea..df25e5307698 100644 --- a/tests/ut/distributed/ascend_store/test_pool_worker.py +++ b/tests/ut/distributed/ascend_store/test_pool_worker.py @@ -163,6 +163,10 @@ def make_worker( config.model_config.get_total_num_hidden_layers.return_value = sum(pp_partition) config.model_config.compute_hash.return_value = "test-model-config" config.cache_config.compute_hash.return_value = "test-cache-config" + # Exercise main's native hash protocol while retaining the release + # fixture's mutable identity used by the topology-isolation assertions. + config.model_config.compile_factors.side_effect = lambda: {"identity": config.model_config.compute_hash()} + config.cache_config.compile_factors.side_effect = lambda: {"identity": config.cache_config.compute_hash()} config.cache_config.cache_dtype = "auto" config.speculative_config = None config.additional_config = {"enable_kvpp": use_kvpp} diff --git a/tests/ut/distributed/eplb/test_state.py b/tests/ut/distributed/eplb/test_state.py index b72cf9136ff9..8d5550b75372 100644 --- a/tests/ut/distributed/eplb/test_state.py +++ b/tests/ut/distributed/eplb/test_state.py @@ -19,6 +19,7 @@ AscendEplbState, ) from vllm_ascend.distributed.eplb.policy.stair import StairEplbPolicy +from vllm_ascend.utils import get_config_hash def test_update_mapping_refreshes_ascend_routing_without_replacing_state(): @@ -28,8 +29,10 @@ def test_update_mapping_refreshes_ascend_routing_without_replacing_state(): state._suspended = False model_config = MagicMock() model_config.compute_hash.return_value = "target" + model_config.compile_factors.return_value = {"identity": "target"} model_state = object() state.model_states = {"target": model_state} + state.model_states[get_config_hash(model_config, cached=True)] = model_state mapping = torch.tensor([[1, 0]], dtype=torch.int32) events = [] with ( @@ -216,6 +219,7 @@ def test_add_model_snapshots_the_placement_on_every_profile(monkeypatch, device) monkeypatch.setattr(upstream_eplb_state.EplbState, "add_model", lambda *_, **__: None) model_config = MagicMock() model_config.compute_hash.return_value = "target" + model_config.compile_factors.return_value = {"identity": "target"} physical = torch.tensor([[0, 1, 0, 1]]) state = AscendEplbState.__new__(AscendEplbState) @@ -223,6 +227,7 @@ def test_add_model_snapshots_the_placement_on_every_profile(monkeypatch, device) state._configured_policy = None state.parallel_config = SimpleNamespace(eplb_config=SimpleNamespace(policy="default")) state.model_states = {"target": SimpleNamespace(physical_to_logical_map=physical)} + state.model_states[get_config_hash(model_config, cached=True)] = state.model_states["target"] monkeypatch.setattr(AscendEplbState, "uses_custom_load_stats", property(lambda _: False)) state.add_model(object(), model_config) diff --git a/tests/ut/distributed/kv_transfer/kv_p2p/test_remote_prefill_lifecycle.py b/tests/ut/distributed/kv_transfer/kv_p2p/test_remote_prefill_lifecycle.py index b07a96eac823..3dab13017f4b 100644 --- a/tests/ut/distributed/kv_transfer/kv_p2p/test_remote_prefill_lifecycle.py +++ b/tests/ut/distributed/kv_transfer/kv_p2p/test_remote_prefill_lifecycle.py @@ -71,7 +71,7 @@ def test_basic_lifecycle(): assert len(block_pool.cached_block_hash_to_block) == 0 blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[request_id] for block in blocks: - assert block._block_hash is None + assert block.block_hash is None # (1b): forward() model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT @@ -104,7 +104,7 @@ def test_basic_lifecycle(): blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[request_id] for block in blocks: assert block.ref_cnt == 1 - num_hashed_blocks += 1 if block._block_hash is not None else 0 + num_hashed_blocks += 1 if block.block_hash is not None else 0 assert num_hashed_blocks == NUM_EXTERNAL_FULL_BLOCKS scheduled_req = scheduler_output.scheduled_new_reqs[0] diff --git a/tests/ut/distributed/weight_transfer/test_hccl_engine.py b/tests/ut/distributed/weight_transfer/test_hccl_engine.py index 84ef5fd4e71e..0ce4f01c3feb 100644 --- a/tests/ut/distributed/weight_transfer/test_hccl_engine.py +++ b/tests/ut/distributed/weight_transfer/test_hccl_engine.py @@ -30,7 +30,10 @@ def _npu_context(): with ( patch.object( - torch, "npu", SimpleNamespace(device=lambda _: nullcontext(), current_stream=MagicMock()), create=True + torch, + "npu", + SimpleNamespace(device=lambda _: nullcontext(), current_stream=MagicMock(), empty_cache=MagicMock()), + create=True, ), patch("torch.accelerator.current_device_index", return_value=0), ): @@ -43,13 +46,18 @@ def _make_engine(): engine.model_config = MagicMock() engine.device = torch.device("cpu") engine._last_hccl_update_info = {} + # Match the upstream constructor's disabled load-plan state. + engine._plan_metadata = None + engine._plan_manifest = None + engine._plan = None return engine def test_start_weight_update_initializes_layerwise_reload(): engine = _make_engine() - with patch("vllm.model_executor.model_loader.reload.initialize_layerwise_reload") as initialize: + start = "initialize_layerwise_reload" if vllm_version_is("0.33.0") else "start_reload" + with patch(f"vllm.model_executor.model_loader.reload.{start}") as initialize: engine.start_weight_update() initialize.assert_called_once_with(engine.model) @@ -58,10 +66,14 @@ def test_start_weight_update_initializes_layerwise_reload(): def test_finish_weight_update_finalizes_layerwise_reload(): engine = _make_engine() - with patch("vllm.model_executor.model_loader.reload.finalize_layerwise_reload") as finalize: + finish = "finalize_layerwise_reload" if vllm_version_is("0.33.0") else "finish_reload" + with patch(f"vllm.model_executor.model_loader.reload.{finish}") as finalize: engine.finish_weight_update() - finalize.assert_called_once_with(engine.model, engine.model_config) + args: tuple[Any, ...] = (engine.model, engine.model_config) + if not vllm_version_is("0.33.0"): + args += (engine.device,) + finalize.assert_called_once_with(*args) def test_shutdown_closes_session_once(): @@ -109,7 +121,10 @@ def test_init_info_preserves_upstream_payload_contract(): assert isinstance(init_info, NCCLWeightTransferInitInfo) assert HCCLWeightTransferEngine.init_info_cls is HCCLWeightTransferInitInfo - assert worker_init_payload(init_info) == dict(master_address="host", master_port=1234, rank_offset=1, world_size=2) + expected = dict(master_address="host", master_port=1234, rank_offset=1, world_size=2) + if not vllm_version_is("0.33.0"): + expected["use_load_plan"] = False + assert worker_init_payload(init_info) == expected restored = HCCLWeightTransferInitInfo(**worker_init_payload(init_info)) assert restored == init_info assert restored.packed is None diff --git a/tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py b/tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py index 4eb52ff02cb2..20f9f8495f1c 100644 --- a/tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py +++ b/tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py @@ -40,7 +40,10 @@ import sys import types import weakref +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import replace +from importlib import import_module from unittest.mock import MagicMock, patch import pytest @@ -52,6 +55,7 @@ NPUIPCTrainerWeightTransferEngine, NPUIPCWeightTransferEngine, ) +from vllm_ascend.utils import vllm_version_is _MODULE = "vllm_ascend.distributed.weight_transfer.npu_ipc_engine" @@ -78,6 +82,8 @@ def _patch_reload_module(*, initialize=None, finalize=None): fake_mod = types.ModuleType("vllm.model_executor.model_loader.reload") fake_mod.initialize_layerwise_reload = initialize or MagicMock() # type: ignore[attr-defined] fake_mod.finalize_layerwise_reload = finalize or MagicMock() # type: ignore[attr-defined] + fake_mod.start_reload = fake_mod.initialize_layerwise_reload # type: ignore[attr-defined] + fake_mod.finish_reload = fake_mod.finalize_layerwise_reload # type: ignore[attr-defined] return patch.dict( sys.modules, {"vllm.model_executor.model_loader.reload": fake_mod}, @@ -198,13 +204,17 @@ def test_finish_weight_update(): engine = object.__new__(NPUIPCWeightTransferEngine) engine.model = MagicMock() engine.model_config = MagicMock() + engine.device = torch.device("cpu") engine._packed_importer = MagicMock() mock_finalize = MagicMock() with _patch_reload_module(finalize=mock_finalize): engine.finish_weight_update() - mock_finalize.assert_called_once_with(engine.model, engine.model_config) + if vllm_version_is("0.33.0"): + mock_finalize.assert_called_once_with(engine.model, engine.model_config) + else: + mock_finalize.assert_called_once_with(engine.model, engine.model_config, engine.device) engine._packed_importer.close.assert_called_once_with() @@ -342,8 +352,33 @@ def test_packed_receive_rejects_missing_sizes(): engine.model.load_weights.assert_not_called() +def _legacy_mtp_validation_fixture(): + """Supply the removed release API when testing its engine branch on main.""" + enabled = ContextVar("test_mtp_completeness_check_enabled", default=True) + + @contextmanager + def disable_mtp_completeness_check(): + token = enabled.set(False) + try: + yield + finally: + enabled.reset(token) + + return types.SimpleNamespace( + is_mtp_completeness_check_enabled=enabled.get, + disable_mtp_completeness_check=disable_mtp_completeness_check, + ) + + def test_receive_scopes_mtp_completeness_check_to_loading(): - from vllm.model_executor.model_loader.mtp_validation import is_mtp_completeness_check_enabled + # This test targets the release-only engine branch on both CPU lanes. + # The main branch (no legacy import) is covered by Nov10 contract tests. + mtp_validation = ( + import_module("vllm.model_executor.model_loader.mtp_validation") + if vllm_version_is("0.33.0") + else _legacy_mtp_validation_fixture() + ) + is_mtp_completeness_check_enabled = mtp_validation.is_mtp_completeness_check_enabled engine = object.__new__(NPUIPCWeightTransferEngine) engine.device = types.SimpleNamespace(index=1) @@ -359,6 +394,8 @@ def load(weights): engine.model.load_weights.side_effect = load with ( + patch(f"{_MODULE}.vllm_version_is", return_value=True), + patch(f"{_MODULE}.import_module", return_value=mtp_validation), _patch_rebuild_npu_tensor(lambda *args: torch.ones(1)), patch(f"{_MODULE}.npu_generate_uuid", return_value="node-1"), pytest.raises(RuntimeError, match="load failed"), diff --git a/tests/ut/graph_dispatch_fixture.py b/tests/ut/graph_dispatch_fixture.py new file mode 100644 index 000000000000..dd42121bd500 --- /dev/null +++ b/tests/ut/graph_dispatch_fixture.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Keep platform graph-selection assertions on the active runner's real API. + +This CPU fixture adapts the old key-initialization/tuple interface to MRv2's +descriptor interface. It builds real native dispatch tables; it does not +capture or replay any device graph. +""" + +from importlib import import_module +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager + +from vllm_ascend.utils import vllm_version_is + + +class _ModelRunnerV2Dispatcher: + def __init__(self, vllm_config): + self.vllm_config = vllm_config + + def initialize_cudagraph_keys(self, mode): + query_len = self.vllm_config.uniform_decode_query_len + pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True) + with ( + patch("vllm.v1.worker.gpu.cudagraph_utils.get_pp_group", return_value=pp_group), + patch("vllm.v1.worker.gpu.cudagraph_utils.current_platform.get_global_graph_pool", return_value=None), + ): + self.manager = CudaGraphManager( + self.vllm_config, + torch.device("cpu"), + mode, + {query_len: self.vllm_config.scheduler_config.max_num_seqs}, + decode_request_sizes=[ + tokens // query_len + for tokens in self.vllm_config.compilation_config.cudagraph_capture_sizes + if tokens % query_len == 0 + ], + ) + # dispatch() only selects initialized descriptors after this gate. + # No graph is captured, installed in manager.graphs, or replayed here. + self.manager._graphs_captured = True + + def dispatch(self, num_tokens, uniform_decode=False): + query_len = self.vllm_config.uniform_decode_query_len + desc = self.manager.dispatch( + num_reqs=num_tokens // query_len if uniform_decode else 1, + num_tokens=num_tokens, + uniform_token_count=query_len if uniform_decode else None, + num_active_loras=0, + ) + return desc.cg_mode, desc + + +if vllm_version_is("0.33.0"): + CudagraphDispatcher = import_module("vllm.v1.cudagraph_dispatcher").CudagraphDispatcher +else: + CudagraphDispatcher = _ModelRunnerV2Dispatcher diff --git a/tests/ut/kv_offload/test_dspark_transfer_cancellation.py b/tests/ut/kv_offload/test_dspark_transfer_cancellation.py index 975ef5d22819..ec1f9bd8c8ac 100644 --- a/tests/ut/kv_offload/test_dspark_transfer_cancellation.py +++ b/tests/ut/kv_offload/test_dspark_transfer_cancellation.py @@ -147,12 +147,16 @@ def _fixture(names=("cancel",), *, admit=True): upstream.waiting = Mock() upstream.skipped_waiting = Mock() upstream.kv_holding_waiting = Mock() + upstream.async_kv_waiting = Mock() + upstream._aborted_kv_loads = set() upstream.deferred_waiting = set() upstream.finished_recving_kv_req_ids = set() upstream.failed_recving_kv_req_ids = set() upstream.finished_req_ids = set() upstream.finished_req_ids_dict = None upstream._inflight_prefills = set() + upstream._parked_reservations = {} + upstream._parked_reserved_blocks = 0 upstream._kv_fetch_stages = None upstream.aux_output_connector = None upstream.ec_connector = None diff --git a/tests/ut/kv_offload/utils.py b/tests/ut/kv_offload/utils.py index b9f55fdb71f9..e10ad8f0c03f 100644 --- a/tests/ut/kv_offload/utils.py +++ b/tests/ut/kv_offload/utils.py @@ -29,6 +29,8 @@ from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager +from vllm_ascend.utils import vllm_version_is + EOS_TOKEN_ID = 50256 @@ -59,6 +61,7 @@ def _fake_opt_model_info() -> SimpleNamespace: """Static ModelInfo for OPT fake weights; avoids registry subprocess inspect.""" return SimpleNamespace( architecture="OPTForCausalLM", + speculator_config_class=None, is_text_generation_model=True, is_pooling_model=False, attn_type="decoder", @@ -257,7 +260,7 @@ def create_model_runner_output( model_runner_output = ModelRunnerOutput( req_ids=req_ids, - req_id_to_index=req_id_to_index, + **({"req_id_to_index": req_id_to_index} if vllm_version_is("0.33.0") else {}), sampled_token_ids=sampled_token_ids, logprobs=None, prompt_logprobs_dict={}, diff --git a/tests/ut/models/minimax_m3/test_minimax_m3.py b/tests/ut/models/minimax_m3/test_minimax_m3.py index 8a6bdacc3671..983fd050cd5b 100644 --- a/tests/ut/models/minimax_m3/test_minimax_m3.py +++ b/tests/ut/models/minimax_m3/test_minimax_m3.py @@ -128,6 +128,57 @@ def _make_decoder_layer(*, use_sequence_parallel: bool, is_moe: bool = True) -> class TestMiniMaxM3Modeling(unittest.TestCase): + def test_pp_auxiliary_states_match_unsplit_target(self): + from vllm.v1.worker.gpu.pp_utils import PPHandler + + from vllm_ascend.utils import vllm_version_is + + class TargetLayer(nn.Module): + use_sequence_parallel = False + + def forward(self, positions, hidden, residual, input_is_sequence_parallel=False): + return hidden + 1, hidden * 10 + + module = "vllm_ascend.models.minimax_m3.minimax_m3" + + def run(cuts): + incoming = None + for stage, (start, end) in enumerate(zip(cuts, cuts[1:])): + group = SimpleNamespace( + world_size=len(cuts) - 1, is_first_rank=stage == 0, is_last_rank=stage == len(cuts) - 2 + ) + model = MiniMaxM3Model.__new__(MiniMaxM3Model) + nn.Module.__init__(model) + model.start_layer, model.end_layer = start, end + model.layers = nn.ModuleList([TargetLayer() for _ in range(6)]) + model._enable_eagle3_aux_hidden_states = True + model.embed_input_ids = lambda ids: torch.ones(2, 3) + model.norm = lambda hidden, residual: (hidden + residual, None) + with ( + patch(f"{module}.get_pp_group", return_value=group), + patch(f"{module}.is_forward_context_available", return_value=False), + patch("vllm.distributed.parallel_state.get_pp_group", return_value=group), + patch("vllm.distributed.parallel_state.model_parallel_is_initialized", return_value=True), + ): + model.set_aux_hidden_state_layers((0, 1, 2, 4, 6)) + handler = SimpleNamespace() + PPHandler.configure_aux_hidden_state_relay(handler, SimpleNamespace(model=model)) + output = model.forward( + torch.zeros(2, dtype=torch.long) if group.is_first_rank else None, torch.arange(2), incoming + ) + if not group.is_last_rank: + incoming = PPHandler.relay_aux_hidden_states(handler, incoming, output) + elif vllm_version_is("0.33.0"): + return output + else: + return PPHandler.add_upstream_aux_hidden_states(handler, incoming, output) + + expected = run((0, 6)) + for cuts in ((0, 2, 4, 6), (0, 0, 3, 6), (0, 3, 6, 6)): + actual = run(cuts) + torch.testing.assert_close(actual, expected) + self.assertEqual(len(actual[1]), 5) + def test_dspark_receives_residual_added_auxiliary_states(self): class _TargetLayer(nn.Module): use_sequence_parallel = False diff --git a/tests/ut/models/minimax_m3/test_msa_m3.py b/tests/ut/models/minimax_m3/test_msa_m3.py index 079250747e6e..609ac786e3cf 100644 --- a/tests/ut/models/minimax_m3/test_msa_m3.py +++ b/tests/ut/models/minimax_m3/test_msa_m3.py @@ -218,7 +218,9 @@ def _make_sparse_builder( device: torch.device, kv_cache_dtype: torch.dtype = torch.bfloat16 ) -> AscendMiniMaxM3SparseMetadataBuilder: vllm_config = _make_vllm_config() - vllm_config.speculative_config = SimpleNamespace(num_speculative_tokens=7, parallel_drafting=False) + vllm_config.speculative_config = SimpleNamespace( + num_speculative_tokens=7, max_num_new_slots_for_drafting=0, parallel_drafting=False + ) spec = FullAttentionSpec( block_size=128, num_kv_heads=2, @@ -2256,7 +2258,9 @@ def test_sparse_builder_keeps_unverified_gbsa_paths_in_layers(dtype, block_size, config = _make_vllm_config() config.model_config.dtype = dtype spec = FullAttentionSpec(block_size=block_size, num_kv_heads=2, head_size=128, dtype=dtype) - config.speculative_config = SimpleNamespace(num_speculative_tokens=3, parallel_drafting=False) + config.speculative_config = SimpleNamespace( + num_speculative_tokens=3, max_num_new_slots_for_drafting=0, parallel_drafting=False + ) with patch.object( msa_m3_module, "get_current_hardware_profile", return_value=get_hardware_profile(AscendDeviceType.A3) ): @@ -2819,7 +2823,9 @@ def test_gbsa_schedule_lengths_follow_speculative_config(spec_tokens, mock_gbsa_ ) config = _make_vllm_config(max_num_batched_tokens=1024) config.speculative_config = ( - SimpleNamespace(num_speculative_tokens=spec_tokens, parallel_drafting=False) if spec_tokens else None + SimpleNamespace(num_speculative_tokens=spec_tokens, max_num_new_slots_for_drafting=0, parallel_drafting=False) + if spec_tokens + else None ) spec = FullAttentionSpec(block_size=128, num_kv_heads=2, head_size=128, dtype=torch.bfloat16) builder = AscendMiniMaxM3SparseMetadataBuilder(spec, ["layer0.attn"], config, torch.device("cpu")) diff --git a/tests/ut/models/test_kimi_k3_adapter.py b/tests/ut/models/test_kimi_k3_adapter.py index 4cccc4ba4050..5e5964a4f6de 100644 --- a/tests/ut/models/test_kimi_k3_adapter.py +++ b/tests/ut/models/test_kimi_k3_adapter.py @@ -418,6 +418,7 @@ def test_kimi_mixed_kda_gate_weights_use_upstream_packed_loader(monkeypatch): model.config = SimpleNamespace( is_moe=False, is_linear_attn=True, + dsa_config=None, linear_attn_config={}, q_lora_rank=None, num_hidden_layers=1, diff --git a/tests/ut/models/test_kimi_k3_eplb.py b/tests/ut/models/test_kimi_k3_eplb.py index 7b23cd272502..2039662c3bcd 100644 --- a/tests/ut/models/test_kimi_k3_eplb.py +++ b/tests/ut/models/test_kimi_k3_eplb.py @@ -313,6 +313,7 @@ def test_expert_loading_initializes_every_replica( num_experts=num_experts, num_hidden_layers=2, num_nextn_predict_layers=1, + dsa_config=None, linear_attn_config={}, q_lora_rank=None, is_linear_attn=False, diff --git a/tests/ut/models/test_kimi_k3_gate_loading.py b/tests/ut/models/test_kimi_k3_gate_loading.py index f436cc527daa..013f3ba78c68 100644 --- a/tests/ut/models/test_kimi_k3_gate_loading.py +++ b/tests/ut/models/test_kimi_k3_gate_loading.py @@ -21,6 +21,8 @@ def make_linear(cls, output_sizes, tp_size=1, tp_rank=0): layer.tp_rank = tp_rank layer.tp_size = tp_size layer.output_sizes = output_sizes + layer.shard_ids = tuple(range(len(output_sizes))) + layer.loaded_shard_ids = {} return layer @@ -38,6 +40,7 @@ def make_model(tp_size=1, tp_rank=0, fused_gate=False): nn.Module.__init__(model) model.config = SimpleNamespace( is_moe=False, + dsa_config=None, is_linear_attn=True, linear_attn_config={"use_full_rank_gate": True}, q_lora_rank=1536, diff --git a/tests/ut/models/test_kimi_k3_pp_sp.py b/tests/ut/models/test_kimi_k3_pp_sp.py index 5a7dc380b3ca..cf4750aca27d 100644 --- a/tests/ut/models/test_kimi_k3_pp_sp.py +++ b/tests/ut/models/test_kimi_k3_pp_sp.py @@ -224,6 +224,9 @@ def runtime(monkeypatch): "get_pp_group": lambda: context.pp, "get_tensor_model_parallel_world_size": lambda: context.tp, "get_tensor_model_parallel_rank": lambda: context.rank, + # BaseModel above models the release's cumulative aux protocol. + # Native main's centralized relay is covered separately with its real mixin/PPHandler. + "vllm_version_is": lambda _: True, } # Transport and sharding use the repository implementations, including # multi-dimensional residuals and zero-length token tensors. @@ -332,6 +335,90 @@ def forward(self, *, hidden_states, positions): return model +@pytest.fixture +def native_aux_runtime(runtime, monkeypatch): + # Keep the standalone CPU fixture importable without a vLLM installation; + # this additional fixture explicitly exercises the installed native ABI. + import vllm.distributed.parallel_state as parallel_state + from vllm.model_executor.models.interfaces import EagleModelMixin + from vllm.v1.worker.gpu.pp_utils import PPHandler + + from vllm_ascend.utils import vllm_version_is + + namespace, context = runtime + release = vllm_version_is("0.33.0") + + class NativeBase(BaseModel, EagleModelMixin): + _cache_aux_pp_layout = EagleModelMixin._cache_aux_pp_layout + _set_aux_hidden_state_layers = EagleModelMixin._set_aux_hidden_state_layers + pack_local_aux_hidden_states = EagleModelMixin.pack_local_aux_hidden_states + AUX_HIDDEN_STATE_KEY = "pp_transport_aux_hidden_states_" + + if not release: + NativeBase.aux_hidden_state_slot_keys = EagleModelMixin.aux_hidden_state_slot_keys + # Alias the original test's legacy diagnostic to native receive slots. + # Production main no longer has or needs this attribute. + setattr( # noqa: B010 - install the main-only diagnostic descriptor on the test class + NativeBase, + "_aux_upstream_total_cached", + property(lambda self: self._aux_slot_base_cached if context.pp.is_last_rank else 0), + ) + namespace.update(BaseModel=NativeBase, vllm_version_is=vllm_version_is) + load_definitions( + "vllm_ascend/models/kimi_k3/kimi_k3.py", + {"AscendKimiLinearModel"}, + namespace, + bases={"AscendKimiLinearModel": "BaseModel"}, + methods={ + "AscendKimiLinearModel": { + "__init__", + "forward", + "_cache_aux_pp_layout", + "make_empty_intermediate_tensors", + "aux_hidden_state_slot_keys", + } + }, + ) + + def pp_group(): + if not hasattr(context.pp, "world_size"): + context.pp.world_size = 3 + return context.pp + + namespace["get_pp_group"] = pp_group + monkeypatch.setattr(parallel_state, "model_parallel_is_initialized", lambda: True) + monkeypatch.setattr(parallel_state, "get_pp_group", pp_group) + model_cls = namespace["AscendKimiLinearModel"] + forward = model_cls.forward + + def forward_with_runner(self, input_ids, positions, intermediate_tensors, inputs_embeds=None): + self._cache_aux_pp_layout() + handler = SimpleNamespace() + if release: + handler.aux_hidden_state_relay_keys = tuple( + f"{self.AUX_HIDDEN_STATE_KEY}{i}" for i in range(self._aux_slot_base_cached) + ) + else: + PPHandler.configure_aux_hidden_state_relay(handler, SimpleNamespace(model=self)) + output = forward(self, input_ids, positions, intermediate_tensors, inputs_embeds) + if not context.pp.is_last_rank: + return PPHandler.relay_aux_hidden_states(handler, intermediate_tensors, output) + if not release: + return PPHandler.add_upstream_aux_hidden_states(handler, intermediate_tensors, output) + return output + + model_cls.forward = forward_with_runner + return namespace, context + + +@pytest.mark.parametrize("materialized", [False, True]) +@pytest.mark.parametrize("sp", [False, True]) +@pytest.mark.parametrize("tp,num_tokens", [(2, 3), (16, 8)]) +@pytest.mark.parametrize("cuts", [(0, 1, 3, 5), (0, 0, 3, 5), (0, 2, 2, 5), (0, 2, 5, 5)]) +def test_native_aux_relay_preserves_pipeline_shards(native_aux_runtime, materialized, sp, tp, num_tokens, cuts): + test_pipeline_shards_match_unsplit_model(native_aux_runtime, materialized, sp, tp, num_tokens, cuts) + + @pytest.mark.parametrize("materialized", [False, True]) @pytest.mark.parametrize("sp", [False, True]) @pytest.mark.parametrize("tp,num_tokens", [(2, 8), (2, 3), (16, 8)]) diff --git a/tests/ut/ops/test_fused_moe.py b/tests/ut/ops/test_fused_moe.py index 62cb674070d2..ef1f1427cfd6 100644 --- a/tests/ut/ops/test_fused_moe.py +++ b/tests/ut/ops/test_fused_moe.py @@ -51,6 +51,7 @@ from vllm_ascend.quantization.methods.base import PreparedLinearInput from vllm_ascend.quantization.methods.w8a8.w8a8_mxfp8 import AscendW8A8MXFP8DynamicLinearMethod from vllm_ascend.quantization.quant_type import QuantType +from vllm_ascend.utils import vllm_version_is def test_prepared_linear_input_import_does_not_cycle_through_moe(): @@ -319,7 +320,6 @@ def test_process_weights_after_loading_splits_lists_for_dynamic_eplb(monkeypatch def test_update_expert_map_updates_routed_experts_and_manager(): routed_experts = AscendRoutedExperts.__new__(AscendRoutedExperts) - nn.Module.__init__(routed_experts) manager = SimpleNamespace(_expert_map=torch.tensor([0, 1, 2, -1], dtype=torch.int32)) routed_experts.expert_map_manager = manager routed_experts.ascend_expert_map = torch.tensor([0, -1, 1], dtype=torch.int32) @@ -2106,6 +2106,8 @@ def init_base_runner( routed_input_transform, routed_output_transform, routed_scaling_factor, + divide_shared_output_in_fp16=False, + routed_scale_mode=None, ): nn.Module.__init__(runner) runner.layer_name = layer_name @@ -2116,6 +2118,10 @@ def init_base_runner( runner.routed_input_transform = routed_input_transform runner.routed_output_transform = routed_output_transform runner.routed_scaling_factor = routed_scaling_factor + if routed_scale_mode is not None: + runner.routed_scale_mode = routed_scale_mode + divide_shared_output_in_fp16 = routed_scale_mode.value == "output_fp16_guard" + runner.divide_shared_output_in_fp16 = divide_shared_output_in_fp16 runner._forward_entry = upstream_forward_entry monkeypatch.setattr(fused_moe_module.MoERunner, "__init__", init_base_runner) @@ -3019,7 +3025,15 @@ def test_forward_impl_keeps_full_width_input_for_shared_experts(monkeypatch): assert result[1] is routed_out -def _stub_moe_runner_init(monkeypatch, *, gate=None, shared_experts=None, fused_mc2_comm=None): +def _stub_moe_runner_init( + monkeypatch, + *, + gate=None, + shared_experts=None, + fused_mc2_comm=None, + divide_shared_output_in_fp16=False, + routed_scale_mode=None, +): """Construct AscendMoERunner with a lightweight MoERunner.__init__ stub.""" moe_config = SimpleNamespace(hidden_dim=4, ep_size=1) routed_experts = SimpleNamespace( @@ -3042,6 +3056,8 @@ def init_base_runner( routed_input_transform, routed_output_transform, routed_scaling_factor, + divide_shared_output_in_fp16=False, + routed_scale_mode=None, ): nn.Module.__init__(runner) runner.layer_name = layer_name @@ -3053,6 +3069,10 @@ def init_base_runner( runner.routed_input_transform = routed_input_transform runner.routed_output_transform = routed_output_transform runner.routed_scaling_factor = routed_scaling_factor + if routed_scale_mode is not None: + runner.routed_scale_mode = routed_scale_mode + divide_shared_output_in_fp16 = routed_scale_mode.value == "output_fp16_guard" + runner.divide_shared_output_in_fp16 = divide_shared_output_in_fp16 runner._forward_entry = object() monkeypatch.setattr(fused_moe_module.MoERunner, "__init__", init_base_runner) @@ -3078,9 +3098,28 @@ def get_comm_method(kind, config): routed_experts=routed_experts, gate=gate, shared_experts=shared_experts, + divide_shared_output_in_fp16=divide_shared_output_in_fp16, + routed_scale_mode=routed_scale_mode, ) +@pytest.mark.parametrize("divide_shared", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("with_shared", [False, True]) +def test_runner_preserves_native_shared_scaling(monkeypatch, divide_shared, dtype, with_shared): + runner = _stub_moe_runner_init(monkeypatch, divide_shared_output_in_fp16=divide_shared) + runner.routed_scaling_factor = 2.0 + if not vllm_version_is("0.33.0"): + assert runner.divide_shared_output_in_fp16 is divide_shared + shared = torch.full((2, 4), 2.0, dtype=dtype) if with_shared else None + fused = torch.full((2, 4), 3.0, dtype=dtype) + shared, fused = runner._maybe_apply_routed_scale_to_output(shared, fused) + divides = (vllm_version_is("0.33.0") or divide_shared) and with_shared and dtype == torch.float16 + torch.testing.assert_close(fused, torch.full_like(fused, 3.0 if divides else 6.0), rtol=0, atol=0) + if shared is not None: + torch.testing.assert_close(shared, torch.full_like(shared, 1.0 if divides else 2.0), rtol=0, atol=0) + + def test_runner_keeps_mega_moe_activation_with_each_layer(monkeypatch): first_kwargs = {"activation": "situglu", "activation_params": {"beta": 4.0, "linear_beta": 25.0}} first = _stub_moe_runner_init(monkeypatch, fused_mc2_comm=SimpleNamespace(mega_moe_activation_kwargs=first_kwargs)) @@ -3334,3 +3373,26 @@ def forward_entry(x, *args): assert not any("ascend_moe_forward_complete" in str(node.target) for node in graph.graph.nodes) finally: library._destroy() + + +@pytest.mark.parametrize("mode_name", ["ROUTER", "OUTPUT", "OUTPUT_FP16_GUARD"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("with_shared", [False, True]) +def test_runner_preserves_native_routed_scale_mode(monkeypatch, mode_name, dtype, with_shared): + main = not vllm_version_is("0.33.0") + mode = None + if main: + from vllm.model_executor.layers.fused_moe.config import RoutedScaleMode + + mode = RoutedScaleMode[mode_name] + runner = _stub_moe_runner_init(monkeypatch, routed_scale_mode=mode) + runner.routed_scaling_factor = 2.0 + if main: + assert runner.routed_scale_mode is mode + shared = torch.full((2, 4), 2.0, dtype=dtype) if with_shared else None + fused = torch.full((2, 4), 3.0, dtype=dtype) + shared, fused = runner._maybe_apply_routed_scale_to_output(shared, fused) + divides = (not main or mode_name == "OUTPUT_FP16_GUARD") and with_shared and dtype == torch.float16 + torch.testing.assert_close(fused, torch.full_like(fused, 3.0 if divides else 6.0), rtol=0, atol=0) + if shared is not None: + torch.testing.assert_close(shared, torch.full_like(shared, 1.0 if divides else 2.0), rtol=0, atol=0) diff --git a/tests/ut/ops/test_gdn_attn_builder.py b/tests/ut/ops/test_gdn_attn_builder.py index 68d667c877aa..ec16c95a595f 100644 --- a/tests/ut/ops/test_gdn_attn_builder.py +++ b/tests/ut/ops/test_gdn_attn_builder.py @@ -199,6 +199,7 @@ def _make_vllm_config( if num_speculative_tokens > 0: speculative_config = SimpleNamespace( num_speculative_tokens=num_speculative_tokens, + max_num_new_slots_for_drafting=0, parallel_drafting=False, enable_adaptive_verification=enable_adaptive_verification, use_eagle_block_drop=lambda: False, diff --git a/tests/ut/ops/test_linear.py b/tests/ut/ops/test_linear.py index 20576ea45520..14a71796b876 100644 --- a/tests/ut/ops/test_linear.py +++ b/tests/ut/ops/test_linear.py @@ -1,6 +1,7 @@ import unittest from contextlib import contextmanager from types import SimpleNamespace +from typing import Any from unittest import mock from unittest.mock import MagicMock, patch @@ -14,6 +15,7 @@ from vllm_ascend.ops.linear import ( AscendDCPGroupColumnParallelLinear, AscendMergedColumnParallelLinear, + AscendQKVParallelLinear, AscendReplicatedLinear, AscendRowParallelLinear, AscendUnquantizedLinearMethod, @@ -21,6 +23,7 @@ from vllm_ascend.ops.linear_op import DCPGroupColumnParallelOp from vllm_ascend.quantization.method_adapters import AscendLinearMethod from vllm_ascend.quantization.methods.w8a8.w8a8_dynamic import AscendW8A8DynamicLinearMethod +from vllm_ascend.utils import vllm_version_is class BaseLinearTest(unittest.TestCase): @@ -227,6 +230,28 @@ def test_oproj_tp(self, mock_get_current_vllm_config): class TestAscendMergedColumnParallelLinear(BaseLinearTest): + @unittest.skipIf(vllm_version_is("0.33.0"), "Stacked shard tracking was added after v0.33") + def test_custom_allocations_preserve_missing_shard_validation(self): + layers: tuple[Any, ...] = ( + AscendMergedColumnParallelLinear(input_size=8, output_sizes=[4, 4], bias=False, disable_tp=True), + AscendQKVParallelLinear(hidden_size=8, head_size=2, total_num_heads=2, bias=False, disable_tp=True), + ) + for layer in layers: + with self.subTest(layer=type(layer).__name__): + shards = layer.shard_ids + first = shards[0] + size = layer._get_shard_size_mapping(first) + layer.weight_loader(layer.weight, torch.ones(size, 8), first) + self.assertEqual(layer.pop_missing_shards(), [f"weight.{shard}" for shard in shards[1:]]) + # A new complete load must start fresh and cover every shard. + for shard in shards: + size = layer._get_shard_size_mapping(shard) + layer.weight_loader(layer.weight, torch.ones(size, 8), shard) + self.assertEqual(layer.pop_missing_shards(), []) + torch.testing.assert_close(layer.weight, torch.ones_like(layer.weight)) + with self.assertRaises(ValueError): + layer.validate_shard_id("invalid", layer.weight) + def test_merged_mlp_tp_init(self): ascend_config._ASCEND_CONFIG = MagicMock() ascend_config._ASCEND_CONFIG.scheduler_config.recompute_scheduler_enable = False diff --git a/tests/ut/ops/test_mla.py b/tests/ut/ops/test_mla.py index 440774720d24..185f0db64ccd 100644 --- a/tests/ut/ops/test_mla.py +++ b/tests/ut/ops/test_mla.py @@ -393,30 +393,34 @@ def test_initialization(self, mock_tp_size, mock_get_vllm_config, mock_indexer_c mock_mla_attn.impl = MagicMock() mock_mla_attn.impl.process_weights_after_loading = MagicMock() - with patch("vllm_ascend.ops.mla.AscendMLAAttention", return_value=mock_mla_attn): + with patch("vllm_ascend.ops.mla.AscendMLAAttention", return_value=mock_mla_attn) as backend: mock_tp_size.return_value = 2 mock_vllm_config = MagicMock(spec=VllmConfig) mock_vllm_config.model_config.hf_text_config = MagicMock(num_hidden_layers=32, first_k_dense_replace=True) mock_get_vllm_config.return_value = mock_vllm_config - mock_vllm_config.compilation_config = CompilationConfig() - - attn = AscendMultiHeadLatentAttention( - hidden_size=self.hidden_size, - num_heads=self.num_heads, - scale=self.scale, - qk_nope_head_dim=self.qk_nope_head_dim, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.kv_lora_rank, - mla_modules=self.mock_mla_modules, - cache_config=self.mock_cache_config, - quant_config=self.mock_quant_config, - prefix=self.prefix, - ) - - self.assertEqual(attn.tp_size, 2) - self.assertIsNotNone(attn.mla_attn) + # CUDA eligibility must not alter Ascend cache/indexer layout. + for shardq_eligible in (False, True): + mock_vllm_config.compilation_config = CompilationConfig() + + attn = AscendMultiHeadLatentAttention( + hidden_size=self.hidden_size, + num_heads=self.num_heads, + scale=self.scale, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + mla_modules=self.mock_mla_modules, + cache_config=self.mock_cache_config, + quant_config=self.mock_quant_config, + prefix=self.prefix, + shardq_eligible=shardq_eligible, + ) + + self.assertEqual(attn.tp_size, 2) + self.assertIsNotNone(attn.mla_attn) + self.assertNotIn("shardq_eligible", backend.call_args.kwargs) @patch("vllm_ascend.ops.mla.IndexerWrapper") @patch("vllm_ascend.ops.mla.get_current_vllm_config") diff --git a/tests/ut/ops/test_mla_weight_processing.py b/tests/ut/ops/test_mla_weight_processing.py index 313894b48200..8169a39d01f2 100644 --- a/tests/ut/ops/test_mla_weight_processing.py +++ b/tests/ut/ops/test_mla_weight_processing.py @@ -14,6 +14,7 @@ from vllm_ascend.attention.mla_v1 import AscendMLAImpl from vllm_ascend.attention.sfa_v1 import AscendSFAImpl from vllm_ascend.ops.mla import AscendMultiHeadLatentAttention +from vllm_ascend.utils import vllm_version_is def make_attention(impl): @@ -22,6 +23,7 @@ def make_attention(impl): nn.Module.__init__(inner) set_default_quant_scales(inner, register_buffer=True) inner.impl = impl + inner.act_dtype = torch.bfloat16 inner.num_heads = 2 inner.kv_lora_rank = 4 inner.qk_nope_head_dim = 2 @@ -30,6 +32,7 @@ def make_attention(impl): inner.kv_b_proj.quant_method = UnquantizedLinearMethod() inner.is_amx_bmm_enabled = False inner.dcp_q_replicate = False + inner.dcp_manager = None inner.is_aiter_triton_fp4_bmm_enabled = False inner.is_aiter_triton_fp8_bmm_enabled = False inner.quant_config = None @@ -70,14 +73,18 @@ def test_impl_post_load_called_once(is_sfa, act_dtype): # SFA disposes this projection; upstream dense packing must be bypassed. impl.process_weights_after_loading.side_effect = lambda _: setattr(inner.kv_b_proj, "weight", None) - inner.process_weights_after_loading(act_dtype) + inner.act_dtype = act_dtype + if vllm_version_is("0.33.0"): + inner.process_weights_after_loading(act_dtype) + else: + inner.process_weights_after_loading() impl.process_weights_after_loading.assert_called_once_with(act_dtype) if is_sfa: upstream.assert_not_called() assert inner.kv_b_proj.weight is None else: - upstream.assert_called_once_with(act_dtype) + upstream.assert_called_once_with(*([act_dtype] if vllm_version_is("0.33.0") else [])) assert inner.W_UV.shape == (2, 4, 4) assert inner.W_UK_T.shape == (2, 2, 4) @@ -113,7 +120,7 @@ def test_mlapo_post_load_does_not_reprocess_released_weights(is_kv_consumer): ): inner.process_weights_after_loading(torch.bfloat16) - upstream.assert_called_once_with(torch.bfloat16) + upstream.assert_called_once_with(*([torch.bfloat16] if vllm_version_is("0.33.0") else [])) impl.process_weights_after_loading.assert_called_once_with(torch.bfloat16) assert (impl.fused_qkv_a_proj.weight is None) == is_kv_consumer assert (impl.q_proj.weight is None) == is_kv_consumer diff --git a/tests/ut/patch/platform/release_scheduler_fixture.py b/tests/ut/patch/platform/release_scheduler_fixture.py new file mode 100644 index 000000000000..badb63dbd215 --- /dev/null +++ b/tests/ut/patch/platform/release_scheduler_fixture.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exact v0.33.0 cached-request ABI for testing release-only patches on main. + +Copied from vLLM tag v0.33.0 (4ba787ec9959e566d3b587a796dff07673324da5): +v1/core/sched/output.py:CachedRequestData and +v1/core/sched/scheduler.py:Scheduler._make_cached_request_data. +The current main scheduler never receives this fixture or these V1 patches. +Main's native payload is tested separately; this fixture keeps the original +release-patch assertions executable after upstream removes their V1 ABI. +""" + +from __future__ import annotations + +import itertools +from dataclasses import dataclass +from functools import cached_property +from typing import Any + +from vllm.v1.core.kv_cache_manager import KVCacheBlocks +from vllm.v1.request import Request + + +def release_cached_request_abi() -> tuple[type[Any], type[Any]]: + @dataclass + class CachedRequestData: + req_ids: list[str] + # For request ids not in resumed_req_ids, new_block_ids will be appended to + # the request's block IDs. For those in the set, new_block_ids will be used as the + # request's block IDs instead of appending to the existing block IDs. + resumed_req_ids: set[str] + # NOTE(woosuk): new_token_ids is only used for pipeline parallelism. + # When PP is not used, new_token_ids will be empty. + new_token_ids: list[list[int]] + # MRV1-only: For requests not scheduled in the last step, propagate the token ids + # to the connector. Won't contain requests scheduled in the prior step. + all_token_ids: dict[str, list[int]] + new_block_ids: list[tuple[list[int], ...] | None] + num_computed_tokens: list[int] + num_output_tokens: list[int] + + # Version of dataclass repr with token IDs obfuscated. + def anon_repr(self) -> str: + new_token_ids_lens = [len(toks) for toks in self.new_token_ids] + all_token_ids_lens = {req_id: len(toks) for req_id, toks in self.all_token_ids.items()} + return ( + f"CachedRequestData(" + f"req_ids={self.req_ids}," + f"resumed_req_ids={self.resumed_req_ids}," + f"new_token_ids_lens={new_token_ids_lens}," + f"all_token_ids_lens={all_token_ids_lens}," + f"new_block_ids={self.new_block_ids}," + f"num_computed_tokens={self.num_computed_tokens}," + f"num_output_tokens={self.num_output_tokens}" + f")" + ) + + def __repr__(self) -> str: + return self.anon_repr() + + @property + def num_reqs(self) -> int: + return len(self.req_ids) + + @cached_property + def _req_id_to_num_output_tokens(self) -> dict[str, int]: + """Cache mapping of req_id to num_output_tokens for O(1) lookup. + + This cached property is safe because CachedRequestData instances + are created fresh each scheduling iteration and not mutated during + computation of iteration details. + """ + return dict(zip(self.req_ids, self.num_output_tokens)) + + def is_context_phase(self, req_id: str) -> bool: + num_output_tokens = self._req_id_to_num_output_tokens.get(req_id) + return num_output_tokens is not None and num_output_tokens == 0 + + @classmethod + def make_empty(cls) -> CachedRequestData: + return cls( + req_ids=[], + resumed_req_ids=set(), + new_token_ids=[], + all_token_ids={}, + new_block_ids=[], + num_computed_tokens=[], + num_output_tokens=[], + ) + + class Scheduler: + use_pp: bool + scheduler_config: Any + use_v2_model_runner: bool + prev_step_scheduled_req_ids: set[str] + + def _make_cached_request_data( + self, + running_reqs: list[Request], + resumed_reqs: list[Request], + num_scheduled_tokens: dict[str, int], + spec_decode_tokens: dict[str, list[int]], + req_to_new_blocks: dict[str, KVCacheBlocks], + ) -> CachedRequestData: + req_ids: list[str] = [] + new_token_ids: list[list[int]] = [] + new_block_ids: list[tuple[list[int], ...] | None] = [] + all_token_ids: dict[str, list[int]] = {} + num_computed_tokens: list[int] = [] + num_output_tokens: list[int] = [] + resumed_req_ids = set() + + num_running_reqs = len(running_reqs) + for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)): + req_id = req.request_id + req_ids.append(req_id) + # NOTE: In PP+async scheduling, we consume token ids via a direct GPU + # broadcast path (`input_batch.prev_sampled_token_ids`), so we can + # omit this payload. + if self.use_pp and not self.scheduler_config.async_scheduling: + # When using PP, the scheduler sends the sampled tokens back, + # because there's no direct communication between the first- + # stage worker and the last-stage worker. Otherwise, we don't + # need to send the sampled tokens back because the model runner + # will cache them. + num_tokens = num_scheduled_tokens[req_id] - len(spec_decode_tokens.get(req_id, ())) + token_ids = req.all_token_ids[req.num_computed_tokens : req.num_computed_tokens + num_tokens] + new_token_ids.append(token_ids) + if idx >= num_running_reqs: + resumed_req_ids.add(req_id) + if not self.use_v2_model_runner: # noqa: SIM102 + if req_id not in self.prev_step_scheduled_req_ids: + all_token_ids[req_id] = req.all_token_ids.copy() + new_block_ids.append(req_to_new_blocks[req_id].get_block_ids(allow_none=True)) + num_computed_tokens.append(req.num_computed_tokens) + num_output_tokens.append(req.num_output_tokens + req.num_output_placeholders) + + return CachedRequestData( + req_ids=req_ids, + resumed_req_ids=resumed_req_ids, + new_token_ids=new_token_ids, + all_token_ids=all_token_ids, + new_block_ids=new_block_ids, + num_computed_tokens=num_computed_tokens, + num_output_tokens=num_output_tokens, + ) + + return Scheduler, CachedRequestData diff --git a/tests/ut/patch/platform/test_patch_balance_schedule.py b/tests/ut/patch/platform/test_patch_balance_schedule.py index e81a68a802fe..7f0ea8ea0846 100644 --- a/tests/ut/patch/platform/test_patch_balance_schedule.py +++ b/tests/ut/patch/platform/test_patch_balance_schedule.py @@ -100,6 +100,7 @@ BalanceScheduler, _balance_scheduling_enabled, ) +from vllm_ascend.utils import vllm_version_is # noqa: E402 # Importing any vLLM module can activate the Ascend platform plugin before this # test module finishes importing, so a direct ``Scheduler`` import may already @@ -675,7 +676,7 @@ def _make_output(scheduler): req_ids = [req.request_id for req in scheduler.running] return ModelRunnerOutput( req_ids=req_ids, - req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + **({"req_id_to_index": {req_id: i for i, req_id in enumerate(req_ids)}} if vllm_version_is("0.33.0") else {}), sampled_token_ids=[[1000]] * len(req_ids), logprobs=None, prompt_logprobs_dict={}, @@ -1123,7 +1124,7 @@ def _make_step_output(output, sampled): req_ids = list(output.num_scheduled_tokens) return ModelRunnerOutput( req_ids=req_ids, - req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + **({"req_id_to_index": {req_id: i for i, req_id in enumerate(req_ids)}} if vllm_version_is("0.33.0") else {}), sampled_token_ids=[[1000] if sampled[req_id] else [] for req_id in req_ids], logprobs=None, prompt_logprobs_dict={}, diff --git a/tests/ut/patch/platform/test_patch_fused_moe.py b/tests/ut/patch/platform/test_patch_fused_moe.py index 9ca55c66d7d0..9828a949888c 100644 --- a/tests/ut/patch/platform/test_patch_fused_moe.py +++ b/tests/ut/patch/platform/test_patch_fused_moe.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, patch +import pytest import torch from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts @@ -225,3 +227,79 @@ def test_factory_composes_custom_routed_experts_with_ascend_contract(): assert kwargs["routed_experts_args"]["n_shared_experts"] == 2 assert patch_fused_moe._adapt_routed_experts_cls(_CustomRoutedExperts) is adapted_cls assert patch_fused_moe._adapt_routed_experts_cls(AscendRoutedExperts) is AscendRoutedExperts + + +@pytest.mark.parametrize("scale_output", [False, True]) +@pytest.mark.parametrize("guard_shared", [False, True]) +def test_factory_maps_legacy_scaling_to_the_active_native_contract(scale_output, guard_shared): + router = _Router() + router_factory = MagicMock(return_value=router) + original_factory = MagicMock(return_value=SimpleNamespace(router=router)) + config = SimpleNamespace( + eplb_config=SimpleNamespace(dynamic_eplb=False, expert_map_path=None, num_redundant_experts=0) + ) + with ( + patch.object(patch_fused_moe, "_original_FusedMoE", original_factory), + patch.object(patch_fused_moe, "get_ascend_config", return_value=config), + patch.object( + patch_fused_moe, "get_current_vllm_config", return_value=SimpleNamespace(use_v2_model_runner=True) + ), + patch.object(patch_fused_moe, "create_ascend_fused_moe_router", router_factory), + ): + patch_fused_moe._ascend_FusedMoE( + num_experts=8, + top_k=2, + routed_scaling_factor=2.0, + apply_routed_scale_to_output=scale_output, + divide_shared_output_in_fp16=guard_shared, + ) + kwargs = original_factory.call_args.kwargs + if patch_fused_moe.vllm_version_is("0.33.0"): + assert kwargs["apply_routed_scale_to_output"] is scale_output + assert "routed_scale_mode" not in kwargs + assert "divide_shared_output_in_fp16" not in kwargs + scales_output = scale_output + else: + mode = kwargs["routed_scale_mode"] + assert mode.value == ( + "output_fp16_guard" if scale_output and guard_shared else "output" if scale_output else "router" + ) + assert "apply_routed_scale_to_output" not in kwargs + assert "divide_shared_output_in_fp16" not in kwargs + scales_output = mode.scales_output + assert router_factory.call_args.kwargs["routed_scaling_factor"] == (1.0 if scales_output else 2.0) + + +@pytest.mark.parametrize("mode_name", ["ROUTER", "OUTPUT", "OUTPUT_FP16_GUARD"]) +def test_factory_preserves_native_scaling_modes(mode_name): + router = _Router() + router_factory = MagicMock(return_value=router) + original_factory = MagicMock(return_value=SimpleNamespace(router=router)) + config = SimpleNamespace( + eplb_config=SimpleNamespace(dynamic_eplb=False, expert_map_path=None, num_redundant_experts=0) + ) + main = not patch_fused_moe.vllm_version_is("0.33.0") + scaling: dict[str, Any] = {} + if main: + from vllm.model_executor.layers.fused_moe.config import RoutedScaleMode + + scaling["routed_scale_mode"] = RoutedScaleMode[mode_name] + else: + scaling["apply_routed_scale_to_output"] = mode_name != "ROUTER" + with ( + patch.object(patch_fused_moe, "_original_FusedMoE", original_factory), + patch.object(patch_fused_moe, "get_ascend_config", return_value=config), + patch.object( + patch_fused_moe, "get_current_vllm_config", return_value=SimpleNamespace(use_v2_model_runner=True) + ), + patch.object(patch_fused_moe, "create_ascend_fused_moe_router", router_factory), + ): + patch_fused_moe._ascend_FusedMoE(num_experts=8, top_k=2, routed_scaling_factor=2.0, **scaling) + kwargs = original_factory.call_args.kwargs + if main: + assert kwargs["routed_scale_mode"] is scaling["routed_scale_mode"] + assert "apply_routed_scale_to_output" not in kwargs + else: + assert kwargs["apply_routed_scale_to_output"] is scaling["apply_routed_scale_to_output"] + assert "routed_scale_mode" not in kwargs + assert router_factory.call_args.kwargs["routed_scaling_factor"] == (2.0 if mode_name == "ROUTER" else 1.0) diff --git a/tests/ut/patch/platform/test_patch_mamba_config.py b/tests/ut/patch/platform/test_patch_mamba_config.py index 2f9b90461c8a..4f827769535b 100644 --- a/tests/ut/patch/platform/test_patch_mamba_config.py +++ b/tests/ut/patch/platform/test_patch_mamba_config.py @@ -88,6 +88,7 @@ def _config( scheduler_config=SimpleNamespace( disable_hybrid_kv_cache_manager=disable_hybrid, ), + use_flashinfer_replayssm=False, kv_transfer_config=connector, speculative_config=(None if speculative_method is None else SimpleNamespace(method=speculative_method)), cache_config=SimpleNamespace( diff --git a/tests/ut/patch/platform/test_patch_pp_mtp.py b/tests/ut/patch/platform/test_patch_pp_mtp.py index fa3a1de07890..5335f2262e27 100644 --- a/tests/ut/patch/platform/test_patch_pp_mtp.py +++ b/tests/ut/patch/platform/test_patch_pp_mtp.py @@ -8,6 +8,7 @@ from vllm.v1.engine.core import EngineCore from vllm.v1.structured_output import StructuredOutputManager +from tests.ut.patch.platform.release_scheduler_fixture import release_cached_request_abi from vllm_ascend.patch.platform.patch_pp_mtp import ( _apply_patch, _is_pd_prefill_node, @@ -24,6 +25,18 @@ from vllm_ascend.utils import vllm_version_is +@pytest.fixture(autouse=True) +def _release_ipc_payload_abi_on_main(monkeypatch, request): + if vllm_version_is("0.33.0") or not request.node.originalname.startswith("test_pp_ipc_cached_request_data_"): + return + from vllm.v1.core.sched import scheduler as native_module + + legacy_scheduler, _ = release_cached_request_abi() + monkeypatch.setattr(request.module, "Scheduler", legacy_scheduler) + monkeypatch.setattr(native_module, "Scheduler", legacy_scheduler) + _patch_scheduler_make_cached_request_data() + + @pytest.mark.parametrize( "model_type,architecture", [ @@ -328,6 +341,7 @@ def validate_grammar(_tokens): is_finished=lambda: False, use_structured_output=True, all_token_ids=[200], + _all_token_ids=[200], structured_output_request=SimpleNamespace( grammar=SimpleNamespace(validate_tokens=validate_grammar), reasoning_ended=False, @@ -338,7 +352,9 @@ def validate_grammar(_tokens): monkeypatch.setattr( manager, "_get_reasoner", - lambda _request: SimpleNamespace(is_reasoning_end_streaming=lambda _ids, _tokens: False), + lambda _request: SimpleNamespace( + reasoning_end_token_ids=set(), is_reasoning_end_streaming=lambda _ids, _tokens: False + ), ) scheduler = SimpleNamespace(requests={"req-0": request}, structured_output_manager=manager) scheduler_output = SimpleNamespace(num_scheduled_tokens={"req-0": 1}) diff --git a/tests/ut/patch/platform/test_patch_swa_bounded_replay.py b/tests/ut/patch/platform/test_patch_swa_bounded_replay.py index 6b7cd02e2dd6..88a7e5130188 100644 --- a/tests/ut/patch/platform/test_patch_swa_bounded_replay.py +++ b/tests/ut/patch/platform/test_patch_swa_bounded_replay.py @@ -27,6 +27,7 @@ from vllm.v1.core.sched.output import CachedRequestData from vllm.v1.core.sched.scheduler import Scheduler +from tests.ut.patch.platform.release_scheduler_fixture import release_cached_request_abi from vllm_ascend.patch.platform import patch_swa_bounded_replay as replay_patch from vllm_ascend.patch.platform.patch_swa_bounded_replay import ( _EXPECTED_PARAMETERS, @@ -34,6 +35,7 @@ _patch_make_cached_request_data, _upstream_carries_replay_start, ) +from vllm_ascend.utils import vllm_version_is pytestmark = pytest.mark.skipif( not _upstream_carries_replay_start(), @@ -44,6 +46,22 @@ REPLAY_START = 64 +@pytest.fixture(autouse=True) +def _release_patch_abi_on_main(monkeypatch, request): + if vllm_version_is("0.33.0"): + return + # These original tests assert the release-only V1 patch contract. Main + # transports replay_start with NewRequestData and must stay unpatched. + legacy_scheduler, legacy_cached = release_cached_request_abi() + monkeypatch.setattr(request.module, "Scheduler", legacy_scheduler) + monkeypatch.setattr(request.module, "CachedRequestData", legacy_cached) + monkeypatch.setattr(replay_patch, "Scheduler", legacy_scheduler) + monkeypatch.setattr(replay_patch, "CachedRequestData", legacy_cached) + monkeypatch.setattr(legacy_scheduler, "_mark_prefix_replay", lambda: None, raising=False) + _patch_cached_request_data_replay_start() + _patch_make_cached_request_data() + + def test_the_probe_finds_upstreams_rewind(): """The probe is the module's own switch, and it is what the skip above rests on: it has to key on something that arrives with the feature.""" @@ -161,3 +179,13 @@ def test_expected_parameters_match_the_patched_method(): """Guards the guard: the tuple above has to describe the method actually installed, or the check would pass on a signature it no longer recognizes.""" assert tuple(inspect.signature(Scheduler._make_cached_request_data).parameters) == _EXPECTED_PARAMETERS + + +def test_native_main_replay_payload_keeps_the_release_patch_inert(): + from vllm.v1.core.sched.output import NewRequestData + from vllm.v1.core.sched.scheduler import Scheduler as NativeScheduler + + assert "replay_start" in NewRequestData.__dataclass_fields__ + if not vllm_version_is("0.33.0"): + assert "resumed_reqs" not in inspect.signature(NativeScheduler._make_cached_request_data).parameters + assert not getattr(NativeScheduler._make_cached_request_data, "_vllm_ascend_swa_replay_patched", False) diff --git a/tests/ut/patch/platform/test_patch_use_v2_model_runner.py b/tests/ut/patch/platform/test_patch_use_v2_model_runner.py index d7ae44550f59..451874c668df 100644 --- a/tests/ut/patch/platform/test_patch_use_v2_model_runner.py +++ b/tests/ut/patch/platform/test_patch_use_v2_model_runner.py @@ -6,6 +6,7 @@ # Importing the module applies the V2-only runner selection at collection time. from vllm_ascend.patch.platform import patch_use_v2_model_runner # noqa: F401 +from vllm_ascend.utils import vllm_version_is def test_use_v2_model_runner_is_driven_by_ascend_default(): @@ -20,7 +21,11 @@ def test_use_v2_model_runner_is_driven_by_ascend_default(): def test_native_v2_spec_pp_validation(method, enable_sp): config = SimpleNamespace( model_config=None, - speculative_config=SimpleNamespace(method=method, parallel_drafting=False), + speculative_config=SimpleNamespace( + method=method, + parallel_drafting=False, + speculator_config=SimpleNamespace(speculator_class="test.Speculator"), + ), compilation_config=SimpleNamespace(mode=CompilationMode.NONE, pass_config=SimpleNamespace(enable_sp=enable_sp)), parallel_config=SimpleNamespace( pipeline_parallel_size=2, @@ -31,7 +36,10 @@ def test_native_v2_spec_pp_validation(method, enable_sp): ), cache_config=SimpleNamespace(mamba_cache_mode="none"), ) - unsupported = VllmConfig._get_v2_model_runner_unsupported_features(config) + if vllm_version_is("0.33.0"): + unsupported = VllmConfig._get_v2_model_runner_unsupported_features(config) + else: + unsupported = VllmConfig._get_model_runner_unsupported_features(config) # Spec+PP is natively supported; unrelated upstream guards still apply. # vLLM main (#60517) dropped the PassConfig.enable_sp guard, so the V2 # runner no longer rejects sequence parallelism there. diff --git a/tests/ut/patch/worker/test_deepseek_pp_embedding.py b/tests/ut/patch/worker/test_deepseek_pp_embedding.py index 5b0ca34bbf8f..6da43a329f0a 100644 --- a/tests/ut/patch/worker/test_deepseek_pp_embedding.py +++ b/tests/ut/patch/worker/test_deepseek_pp_embedding.py @@ -34,7 +34,13 @@ def _make_target(monkeypatch, pp_size, pp_rank, method, *, existing=False, quant config = SimpleNamespace( model_config=SimpleNamespace(hf_config=hf_config), quant_config=quant_config, - speculative_config=None if method is None else SimpleNamespace(method=method), + speculative_config=None + if method is None + else SimpleNamespace( + method=method, + uses_target_hidden_states=True, + draft_model_config=SimpleNamespace(may_use_target_embed=True), + ), use_v2_model_runner=True, ) initial = torch.nn.Embedding(65, 2) if group.is_first_rank or existing else PPMissingLayer() diff --git a/tests/ut/patch/worker/test_patch_bind_kv_cache.py b/tests/ut/patch/worker/test_patch_bind_kv_cache.py index c737a2930514..8b4e97042d9a 100644 --- a/tests/ut/patch/worker/test_patch_bind_kv_cache.py +++ b/tests/ut/patch/worker/test_patch_bind_kv_cache.py @@ -1,4 +1,5 @@ from typing import cast +from unittest.mock import Mock import pytest import torch @@ -72,3 +73,57 @@ def test_bind_kv_cache_accepts_materialized_mamba_states( assert layer.kv_cache[1].is_contiguous() assert layer.kv_cache[0].data_ptr() == conv_state.data_ptr() assert layer.kv_cache[1].data_ptr() == temporal_state.data_ptr() + + +@pytest.mark.parametrize("legacy_hook", [False, True]) +@pytest.mark.parametrize("runner_binding", [False, True]) +def test_binding_retains_tuple_layout_across_tracker_lifecycle(monkeypatch, legacy_hook, runner_binding): + name = "model.layers.0.self_attn" + layer = _TensorOnlyAttentionLayer() + allocated = (torch.zeros(2, 16, 1, 8), torch.zeros(2, 16, 1, 8)) + context = {name: layer} + groups: list[object] = [] + hook = Mock() + # Model the API version as well as its corresponding tracker hook. + monkeypatch.setattr(patch_bind_kv_cache, "vllm_version_is", lambda version: legacy_hook) + if legacy_hook: + monkeypatch.setattr(patch_bind_kv_cache.utils, "share_replayssm_ring_trackers", hook, raising=False) + else: + monkeypatch.delattr(patch_bind_kv_cache.utils, "share_replayssm_ring_trackers", raising=False) + + runner_caches: list[object] = [] + if runner_binding: + patch_bind_kv_cache.bind_kv_cache({name: allocated}, context, runner_caches, kv_cache_groups=groups) + assert runner_caches[0] is allocated + else: + patch_bind_kv_cache.bind_kv_cache_to_layers({name: allocated}, context, kv_cache_groups=groups) + + assert layer.kv_cache is allocated + assert not layer.bind_called + if legacy_hook: + hook.assert_called_once_with([name], context, groups) + else: + hook.assert_not_called() + + +def test_binding_keeps_scheduler_storage_separate_from_kernel_view(monkeypatch): + monkeypatch.setattr(patch_bind_kv_cache, "vllm_version_is", lambda version: False) + monkeypatch.delattr(patch_bind_kv_cache.utils, "share_replayssm_ring_trackers", raising=False) + name = "model.layers.0.self_attn" + storage = torch.arange(32).reshape(4, 8) + kernel_view = storage.reshape(8, 4) + layer = _TensorOnlyAttentionLayer() + runner_caches: list[torch.Tensor] = [] + + patch_bind_kv_cache.bind_kv_cache( + {name: storage}, {name: layer}, runner_caches, layer_kv_caches={name: kernel_view} + ) + + assert runner_caches[0] is storage + assert layer.kv_cache is kernel_view + bound_cache = cast(torch.Tensor, layer.kv_cache) + assert runner_caches[0].shape == (4, 8) + assert bound_cache.shape == (8, 4) + assert bound_cache.data_ptr() == storage.data_ptr() + storage[0, 0] = 99 + assert bound_cache[0, 0] == 99 diff --git a/tests/ut/patch/worker/test_patch_deepseek_v2.py b/tests/ut/patch/worker/test_patch_deepseek_v2.py index a9ca6b0c819f..2d41fe414dfa 100644 --- a/tests/ut/patch/worker/test_patch_deepseek_v2.py +++ b/tests/ut/patch/worker/test_patch_deepseek_v2.py @@ -7,6 +7,7 @@ import torch from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model from vllm.sequence import IntermediateTensors +from vllm.v1.worker.gpu.pp_utils import PPHandler import vllm_ascend.patch.worker.patch_deepseek_v2 as patch_deepseek_v2 import vllm_ascend.worker.v2.pp_transport as pp_transport @@ -15,6 +16,7 @@ _patched_forward, _should_skip_indexer_init, ) +from vllm_ascend.utils import vllm_version_is from vllm_ascend.worker.v2.pp_transport import ( PPTransportDataType, add_pp_transport_tensors, @@ -30,6 +32,28 @@ def _config(**overrides) -> SimpleNamespace: return SimpleNamespace(**values) +def test_native_causal_lm_factory_uses_registered_decoder(): + from vllm.model_executor.models import deepseek_v2 as native + from vllm.model_executor.models.interfaces import EagleModelMixin + + assert native.DeepseekV2ForCausalLM.model_cls is native.DeepseekV2Model + if not vllm_version_is("0.33.0"): + assert issubclass(native.DeepseekV2ForCausalLM.model_cls, EagleModelMixin) + + +def _aux_handler_fixture(model, last_rank): + if vllm_version_is("0.33.0"): + return SimpleNamespace( + aux_hidden_state_relay_keys=[f"aux_hidden_states_{i}" for i in range(model._aux_slot_base_cached)] + ) + handler = SimpleNamespace() + PPHandler.configure_aux_hidden_state_relay(handler, SimpleNamespace(model=model)) + # Keep the existing count assertion as an alias of main's actual receive + # keys. This diagnostic exists only on the test instance, not production. + model._aux_upstream_total_cached = len(handler.upstream_aux_keys) if last_rank else 0 + return handler + + @pytest.mark.parametrize("non_causal", [None, False, True]) def test_mla_constructor_preserves_decode_mask_flag(monkeypatch, non_causal): for name in ("DeepSeekV2FusedQkvAProjLinear", "ColumnParallelLinear", "RowParallelLinear", "RMSNorm"): @@ -163,6 +187,7 @@ def run(split): monkeypatch.setattr(parallel_state, "model_parallel_is_initialized", lambda: True) monkeypatch.setattr(parallel_state, "get_pp_group", lambda group=group: group) model._set_aux_hidden_state_layers(aux_layers) + handler = _aux_handler_fixture(model, last) # IDs are boundaries after a target layer. The state at a # PP cut belongs to the upstream stage, including boundary38. expected_incoming = 0 if first else sum(idx <= start for idx in aux_layers) @@ -172,17 +197,14 @@ def run(split): if not last: if native: # Use the actual upstream relay, without creating streams or process groups. - handler = SimpleNamespace( - aux_hidden_state_relay_keys=[ - f"aux_hidden_states_{i}" for i in range(model._aux_slot_base_cached) - ] - ) output = PPHandler.relay_aux_hidden_states(handler, incoming, output) prefix = "aux_hidden_states_" if native else "pp_transport_aux_hidden_states_" expected_count = sum(idx <= end for idx in aux_layers) assert set(output.tensors) == {"hidden_states", "residual"} | { f"{prefix}{idx}" for idx in range(expected_count) } + elif native and not vllm_version_is("0.33.0"): + output = PPHandler.add_upstream_aux_hidden_states(handler, incoming, output) incoming = output return output diff --git a/tests/ut/patch/worker/test_patch_dspark_eplb.py b/tests/ut/patch/worker/test_patch_dspark_eplb.py index ef2c6ad9bc81..aee20f7a46e3 100644 --- a/tests/ut/patch/worker/test_patch_dspark_eplb.py +++ b/tests/ut/patch/worker/test_patch_dspark_eplb.py @@ -13,6 +13,7 @@ from vllm_ascend import logger as ascend_logger from vllm_ascend.patch.worker import patch_dspark from vllm_ascend.platform import _validate_eplb_config +from vllm_ascend.utils import vllm_version_is @pytest.mark.parametrize("phase", ["prefill", "decode"]) @@ -45,7 +46,8 @@ def dspark_config(tmp_path, monkeypatch, phase, draft_supports_eplb): monkeypatch.setattr(ascend_config, "_ASCEND_CONFIG", None) monkeypatch.setattr(ascend_config, "_INIT_VLLM_CONFIG", None) monkeypatch.setattr(VllmConfig, "use_v2_model_runner", VllmConfig.use_v2_model_runner) - monkeypatch.setattr(VllmConfig, "_validate_v2_model_runner", VllmConfig._validate_v2_model_runner) + validator = "_validate_v2_model_runner" if vllm_version_is("0.33.0") else "_validate_model_runner" + monkeypatch.setattr(VllmConfig, validator, getattr(VllmConfig, validator)) monkeypatch.setattr(ascend_logger, "configure_ascend_file_logging", lambda: None) monkeypatch.setenv("PYTORCH_NPU_ALLOC_CONF", "expandable_segments:True") # Config reconstruction runs before distributed worker initialization. diff --git a/tests/ut/patch/worker/test_patch_mm_input_norm.py b/tests/ut/patch/worker/test_patch_mm_input_norm.py index 9e6756edfc88..420c6b7441c2 100644 --- a/tests/ut/patch/worker/test_patch_mm_input_norm.py +++ b/tests/ut/patch/worker/test_patch_mm_input_norm.py @@ -37,7 +37,7 @@ def test_legacy_shard_loads_and_missing_parameter_still_fails(norm, tmp_path, mo "_render_default_pattern", staticmethod(lambda *args, **kwargs: "model-rank-0-part-*.safetensors"), ) - config = SimpleNamespace(model=str(tmp_path), model_weights=None) + config = SimpleNamespace(model=str(tmp_path), model_weights=None, revision=None) loader.load_weights(model, config) torch.testing.assert_close(model.visual.proj.weight, legacy_state["visual.proj.weight"]) assert set(model.state_dict()) == set(legacy_state) diff --git a/tests/ut/patch/worker/test_patch_v2_gumbel.py b/tests/ut/patch/worker/test_patch_v2_gumbel.py index 71fe5a2a85b6..f0713e882516 100644 --- a/tests/ut/patch/worker/test_patch_v2_gumbel.py +++ b/tests/ut/patch/worker/test_patch_v2_gumbel.py @@ -116,3 +116,40 @@ def test_dspark_sample_logits_dispatch(monkeypatch, probabilistic): result, logits.argmax(dim=-1) + 10, ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_categorical_sampler_applies_markov_bias_before_cache_and_temperature(monkeypatch, dtype, with_bias): + module = importlib.import_module("vllm_ascend.ops.triton.sample.categorical_sample") + logits = torch.tensor([[1.25, -2.5, 4.0], [-3.0, 2.25, 0.5]], dtype=dtype) + original = logits.clone() + bias = torch.tensor([[0.01, 2.0, -6.0], [4.0, -0.1, 1.5]], dtype=dtype) if with_bias else None + cache = torch.empty(2, 2, 3, dtype=dtype) + sampled = torch.tensor([1, 0]) + kernel = MagicMock(return_value=sampled) + monkeypatch.setattr(module, "_categorical_sample", kernel) + result = module.categorical_sample( + logits, + torch.tensor([1, 0]), + torch.tensor([0.5, 1.5]), + torch.tensor([3, 7]), + torch.tensor([8, 12]), + apply_temperature=True, + is_drafting=True, + logits_cache=cache, + bias=bias, + ) + assert result is sampled + kernel.assert_called_once() + passed_logits = kernel.call_args.args[0] + if bias is None: + assert passed_logits is logits + else: + expected = (original.float() + bias.float()).to(dtype) + torch.testing.assert_close(passed_logits, expected, rtol=0, atol=0) + assert passed_logits is not logits + torch.testing.assert_close(logits, original, rtol=0, atol=0) + assert kernel.call_args.args[5] is True + assert kernel.call_args.args[6] is cache + assert kernel.call_args.kwargs["is_drafting"] is True diff --git a/tests/ut/quantization/test_utils.py b/tests/ut/quantization/test_utils.py index 3359bc204358..1a6775784cea 100644 --- a/tests/ut/quantization/test_utils.py +++ b/tests/ut/quantization/test_utils.py @@ -224,7 +224,11 @@ def test_non_quantization_scenarios(self): def test_fa3_quantization_scenario(self): vllm_config = MagicMock() vllm_config.quant_config = AscendModelSlimConfig(FAKQUANT_CONFIG) - vllm_config.kv_transfer_config = KVTransferConfig(kv_connector="MultiConnector", kv_role="kv_consumer") + vllm_config.kv_transfer_config = KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_consumer", + kv_connector_extra_config={"connectors": [{"kv_connector": "MooncakeConnector", "kv_role": "kv_consumer"}]}, + ) vllm_config.cache_config.cache_dtype = "fp8" result = enable_fa_quant(vllm_config) self.assertTrue(result) diff --git a/tests/ut/spec_decode/test_dspark_graph_contract.py b/tests/ut/spec_decode/test_dspark_graph_contract.py index a2e8d3efab49..0cb0f047479c 100644 --- a/tests/ut/spec_decode/test_dspark_graph_contract.py +++ b/tests/ut/spec_decode/test_dspark_graph_contract.py @@ -220,8 +220,14 @@ def test_dispatcher_pads_uniform_draft_descriptors(query_count): manager.compilation_config = SimpleNamespace( cudagraph_capture_sizes=[16, 32, 48, 64, 80, 96, 112, 128], max_cudagraph_capture_size=128, + prefill_cudagraph_capture_token_sizes=[16, 32, 48, 64, 80, 96, 112, 128], + decode_cudagraph_capture_request_sizes=list(range(1, 17)), + max_decode_cudagraph_capture_tokens=None, ) - manager.vllm_config = SimpleNamespace(speculative_config=None) + manager.vllm_config = SimpleNamespace(speculative_config=None, uniform_decode_query_len=query_count) + manager.decode_request_sizes = None + manager.decode_max_reqs = {query_count: 16} + manager._decode_descs = {} manager.max_num_reqs = 16 manager.decode_query_len = query_count manager.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY diff --git a/tests/ut/spec_decode/test_eagle_aclgraph_source_regression.py b/tests/ut/spec_decode/test_eagle_aclgraph_source_regression.py index 6c7a981523cf..e78f83b047ca 100644 --- a/tests/ut/spec_decode/test_eagle_aclgraph_source_regression.py +++ b/tests/ut/spec_decode/test_eagle_aclgraph_source_regression.py @@ -34,6 +34,6 @@ def test_eagle_aclgraph_uses_verified_main_speculator_contract() -> None: def test_eagle_patch_replaces_verified_main_manager_symbol() -> None: source = PATCH.read_text() - assert "vllm_speculator_module.SpeculatorCudaGraphManager = AutoRegressiveAclGraphManager" in source + assert "ar_speculator.SpeculatorCudaGraphManager = AutoRegressiveAclGraphManager" in source assert "PrefillSpeculatorCudaGraphManager" not in source assert "DecodeSpeculatorCudaGraphManager" not in source diff --git a/tests/ut/spec_decode/test_mla_dspark.py b/tests/ut/spec_decode/test_mla_dspark.py index b6f030862831..4c4288ea56ca 100644 --- a/tests/ut/spec_decode/test_mla_dspark.py +++ b/tests/ut/spec_decode/test_mla_dspark.py @@ -4,7 +4,7 @@ from contextlib import contextmanager, nullcontext from types import SimpleNamespace from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np import pytest @@ -247,7 +247,10 @@ def test_raw_prefix_capture_does_not_add_attnres_bank(layer_idx, has_residual): target.config = SimpleNamespace(attn_res_block_size=2) target.aux_hidden_state_layers = (1, 3) - captured = target._maybe_add_hidden_state([], layer_idx, state, residual) + # Isolate allocation while retaining the original raw-prefix/bank checks. + # Native snapshot ownership is exercised without this fixture below. + with patch.object(torch.Tensor, "clone", new=lambda value: value): + captured = target._maybe_add_hidden_state([], layer_idx, state, residual) if layer_idx in target.aux_hidden_state_layers: assert len(captured) == 1 assert captured[0] is state @@ -258,6 +261,22 @@ def test_raw_prefix_capture_does_not_add_attnres_bank(layer_idx, has_residual): torch.testing.assert_close(residual, torch.ones(2, 3, 4)) +def test_raw_prefix_capture_preserves_the_native_snapshot(): + target = AscendKimiLinearModel.__new__(AscendKimiLinearModel) + torch.nn.Module.__init__(target) + target.config = SimpleNamespace(attn_res_block_size=2) + target.aux_hidden_state_layers = (1,) + state = torch.arange(8).view(2, 4) + expected = state.clone() + captured = target._maybe_add_hidden_state([], 1, state, torch.ones(2, 3, 4)) + assert len(captured) == 1 + torch.testing.assert_close(captured[0], expected, rtol=0, atol=0) + if not vllm_version_is("0.33.0"): + assert captured[0] is not state + state.add_(100) + torch.testing.assert_close(captured[0], expected, rtol=0, atol=0) + + def test_padded_mla_query_lengths_are_nested(): spec = make_speculator() metadata = {"draft.0": SimpleNamespace(decode=SimpleNamespace(actual_seq_lengths_q=[5, 5]))} diff --git a/tests/ut/test_ascend_forward_context.py b/tests/ut/test_ascend_forward_context.py index 3e7a253c09f1..1bb43f17596b 100644 --- a/tests/ut/test_ascend_forward_context.py +++ b/tests/ut/test_ascend_forward_context.py @@ -11,6 +11,7 @@ from vllm_ascend.device.hardware import AscendDeviceType from vllm_ascend.device.hardware_profile import get_hardware_profile from vllm_ascend.quantization.quant_type import QuantType +from vllm_ascend.utils import vllm_version_is @pytest.fixture(autouse=True) @@ -63,6 +64,8 @@ def _make_vllm_config( cudagraph_capture_sizes=cudagraph_capture_sizes or [], max_cudagraph_capture_size=max_cudagraph_capture_size, ) + if not vllm_version_is("0.33.0"): + compilation_config.decode_cudagraph_capture_request_sizes = None kv_transfer_config = ( SimpleNamespace(kv_connector=kv_connector, kv_role=kv_role) if kv_connector is not None or kv_role is not None @@ -194,6 +197,49 @@ def test_set_mc2_tokens_capacity_prefill_mc2_uses_max_num_batched_tokens(monkeyp assert afc.get_mc2_tokens_capacity() == 520 +@pytest.mark.parametrize( + "prefill_tokens,decode_sizes,query_lens,token_cap,tp_size,expected", + [ + (16, [1, 2, 3, 4], {4: 1024}, None, 2, 16), + (16, [1, 2, 3, 4], {4: 256}, None, 2, 16), + (8, [1, 4, 16], {4: 16}, None, 2, 64), + (8, [1, 4, 16], {4: 16}, 20, 2, 20), + (8, [1, 4, 16], {1: 16, 4: 3}, None, 2, 16), + (17, [1, 2], {4: 16}, None, 6, 18), + (4096, [1, 2], {4: 16}, None, 2, 1024), + ], +) +def test_mc2_capacity_excludes_independent_draft_capture_grid( + monkeypatch, prefill_tokens, decode_sizes, query_lens, token_cap, tp_size, expected +): + if vllm_version_is("0.33.0"): + pytest.skip("Separate target/draft graph grids are a main-only API") + + monkeypatch.setattr(afc, "use_cann_megamoe", lambda _: False) + vllm_config = _make_vllm_config( + tensor_parallel_size=tp_size, + cudagraph_capture_sizes=[4, 8, 12, 16, 1024], + max_cudagraph_capture_size=1024, + max_num_batched_tokens=8192, + ) + compilation_config = vllm_config.compilation_config + compilation_config.decode_cudagraph_capture_request_sizes = decode_sizes + compilation_config.max_decode_cudagraph_capture_tokens = token_cap + compilation_config.max_prefill_cudagraph_capture_tokens = prefill_tokens + vllm_config.get_decode_cudagraph_query_lens = lambda _: query_lens + original_sizes = compilation_config.cudagraph_capture_sizes.copy() + + afc.set_mc2_tokens_capacity(vllm_config, max_num_reqs=1024, uniform_decode_query_len=4) + + assert afc.get_mc2_tokens_capacity() == expected + assert compilation_config.cudagraph_capture_sizes == original_sizes + assert compilation_config.max_cudagraph_capture_size == 1024 + assert compilation_config.decode_cudagraph_capture_request_sizes == decode_sizes + if expected == 16: + assert afc._select_fused_or_capacity_moe_comm_method(16, vllm_config, expected) == MoECommType.MC2 + assert afc._select_fused_or_capacity_moe_comm_method(1024, vllm_config, expected) == MoECommType.ALLTOALL + + def test_is_decode_only_node_false_without_kv_transfer(): assert afc._is_decode_only_node(_make_vllm_config()) is False diff --git a/tests/ut/test_main2main_nov10_contracts.py b/tests/ut/test_main2main_nov10_contracts.py new file mode 100644 index 000000000000..d8e174eb49e1 --- /dev/null +++ b/tests/ut/test_main2main_nov10_contracts.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Execute release/main adapter boundaries without loading an NPU runtime.""" + +from __future__ import annotations + +import ast +import builtins +import hashlib +import unittest +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +ROOT = Path(__file__).resolve().parents[2] + + +def load_nodes(path, names, namespace): + tree = ast.parse((ROOT / path).read_text(encoding="utf-8")) + body = [n for n in tree.body if isinstance(n, (ast.FunctionDef, ast.ClassDef)) and n.name in names] + module = ast.Module( + body=[ast.ImportFrom(module="__future__", names=[ast.alias("annotations")], level=0), *body], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(ROOT / path), "exec"), namespace) + return namespace + + +class TestNov10Contracts(unittest.TestCase): + def test_platform_bootstrap_does_not_import_vllm_config_from_utils(self): + source = ROOT / "vllm_ascend/utils.py" + tree = ast.parse(source.read_text(encoding="utf-8")) + imports: list[ast.stmt] = [n for n in tree.body if isinstance(n, (ast.Import, ast.ImportFrom))] + + def bootstrap_import(name, *args, **kwargs): + if name == "vllm.config" or name.startswith("vllm.config."): + raise ImportError("vllm.config re-enters current_platform during platform bootstrap") + return MagicMock() + + ns = {"__builtins__": {**vars(builtins), "__import__": bootstrap_import}} + exec(compile(ast.fix_missing_locations(ast.Module(body=imports, type_ignores=[])), str(source), "exec"), ns) + + def test_lookup_protocol_import_selects_its_own_version_constants(self): + source = ROOT / "vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/lookup_key.py" + tree = ast.parse(source.read_text(encoding="utf-8")) + branch = next(n for n in tree.body if isinstance(n, ast.If) and "TYPE_CHECKING or" in ast.unparse(n.test)) + names = ("LOOKUP_MSG", "RESET_MSG", "RESP_ERR", "RESP_OK") + legacy = SimpleNamespace(**{name: f"release-{name}".encode() for name in names}) + main = SimpleNamespace(**{name: f"main-{name}".encode() for name in names}) + for release, checking in ((True, False), (False, False), (True, True)): + loader = MagicMock(return_value=legacy) + imports: list[str] = [] + + def native_import(name, *args, import_log=imports, **kwargs): + import_log.append(name) + self.assertEqual(name, "vllm.distributed.kv_transfer.kv_connector.v1.mooncake_store.protocol") + return main + + ns = { + "TYPE_CHECKING": checking, + "vllm_version_is": lambda _, selected_release=release: selected_release, + "import_module": loader, + "__builtins__": {**vars(builtins), "__import__": native_import}, + } + exec( + compile(ast.fix_missing_locations(ast.Module(body=[branch], type_ignores=[])), str(source), "exec"), ns + ) + expected = legacy if release and not checking else main + self.assertEqual([ns[name] for name in names], [getattr(expected, name) for name in names]) + if release and not checking: + loader.assert_called_once_with("vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol") + self.assertEqual(imports, []) + else: + loader.assert_not_called() + self.assertEqual(len(imports), 1) + + def test_weight_loaders_import_legacy_context_only_on_release_and_restore_on_error(self): + for engine_file, class_name in ( + ("hccl_engine.py", "HCCLWeightTransferEngine"), + ("npu_ipc_engine.py", "NPUIPCWeightTransferEngine"), + ): + tree = ast.parse( + (ROOT / "vllm_ascend/distributed/weight_transfer" / engine_file).read_text(encoding="utf-8") + ) + cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == class_name) + receive = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "receive_weights") + module = ast.Module( + body=[ast.ImportFrom(module="__future__", names=[ast.alias("annotations")], level=0), receive], + type_ignores=[], + ) + for release in (True, False): + with self.subTest(engine=engine_file, release=release): + state = {"enabled": True} + + @contextmanager + def disable_check(scope=state): + scope["enabled"] = False + try: + yield + finally: + scope["enabled"] = True + + loader = MagicMock(return_value=SimpleNamespace(disable_mtp_completeness_check=disable_check)) + weights = [("w", object())] + + def load(received, expected=weights, scope=state, selected_release=release): + self.assertIs(received, expected) + self.assertEqual(scope["enabled"], not selected_release) + raise RuntimeError("load failed") + + params = {"packed": True, "packed_buffer_size_bytes": 16, "packed_num_buffers": 2} + engine = SimpleNamespace( + packed=True, + device=SimpleNamespace(index=1), + model_update_group=object(), + model=MagicMock(load_weights=load), + _wire_params=lambda _, selected=params: selected, + _last_hccl_update_info={}, + _packed_importer=None, + ) + info = SimpleNamespace( + names=["w"], shapes=[[1]], dtype_names=["float32"], ipc_handles={}, tensor_sizes=[4] + ) + ns: dict[str, Any] = { + "vllm_version_is": lambda _, selected_release=release: selected_release, + "nullcontext": nullcontext, + "import_module": loader, + "npu_generate_uuid": lambda _: "node-1", + "packed_npu_ipc_consumer": lambda expected=weights, **_: expected, + "packed_broadcast_consumer": lambda expected=weights, **kwargs: kwargs["post_unpack_func"]( + expected + ), + "torch": SimpleNamespace(npu=SimpleNamespace(device=lambda _: nullcontext())), + } + exec(compile(ast.fix_missing_locations(module), engine_file, "exec"), ns) + with self.assertRaisesRegex(RuntimeError, "load failed"): + ns["receive_weights"](engine, info) + self.assertTrue(state["enabled"]) + if release: + loader.assert_called_once_with("vllm.model_executor.model_loader.mtp_validation") + else: + loader.assert_not_called() + + def test_new_async_kv_limit_rejects_copied_schedulers_before_requests_can_stall(self): + for release in (True, False): + ns = load_nodes( + "vllm_ascend/platform.py", + {"_copied_scheduler_mode", "_validate_copied_scheduler_async_kv_limit"}, + {"vllm_version_is": lambda _, selected_release=release: selected_release}, + ) + validate = ns["_validate_copied_scheduler_async_kv_limit"] + for mode in ("native", "recompute", "dyntra", "profiling", "balance"): + extension = SimpleNamespace( + recompute_scheduler_enable=mode == "recompute", + dyntra_lb_config=SimpleNamespace(enabled=mode == "dyntra"), + profiling_chunk_config=SimpleNamespace(enabled=mode == "profiling"), + enable_balance_scheduling=mode == "balance", + ) + config = SimpleNamespace(scheduler_config=SimpleNamespace()) + ascend_config = SimpleNamespace(scheduler_config=extension) + # Old release configs and existing fixtures do not have the new field. + validate(config, ascend_config) + for limit in (None, 1, 8): + config.scheduler_config.max_async_kv_loads_per_step = limit + with self.subTest(release=release, mode=mode, limit=limit): + if not release and mode != "native" and limit is not None: + with self.assertRaisesRegex(ValueError, "Leave this limit unset"): + validate(config, ascend_config) + else: + validate(config, ascend_config) + self.assertEqual(config.scheduler_config.max_async_kv_loads_per_step, limit) + + def test_compiler_supports_main_registration_and_release_hash(self): + flags = SimpleNamespace(enable_npugraph_ex=True, enable_static_kernel=False, enable_super_kernel=True) + for release in (True, False): + with self.subTest(release=release): + + class Parent: + def __init_subclass__(cls, selected_release=release): + if not selected_release and "compute_hash" in cls.__dict__: + raise TypeError("main rejects the removed compiler hook") + + ns = load_nodes( + "vllm_ascend/compilation/compiler_interface.py", + {"AscendCompiler"}, + { + "CompilerInterface": Parent, + "vllm_version_is": lambda _, selected_release=release: selected_release, + "get_ascend_config": lambda: SimpleNamespace(ascend_compilation_config=flags), + "torch_npu": SimpleNamespace(__version__="test-version"), + "logger": MagicMock(), + }, + ) + compiler = ns["AscendCompiler"]() + config = object() + factors = compiler.compile_factors(config) + self.assertIs(compiler.vllm_config, config) + self.assertEqual( + factors, + { + "torch_npu_version": "test-version", + "enable_npugraph_ex": True, + "enable_static_kernel": False, + "enable_super_kernel": True, + }, + ) + if release: + self.assertEqual( + compiler.compute_hash(config), hashlib.sha256(str(factors).encode()).hexdigest()[:10] + ) + else: + self.assertNotIn("compute_hash", type(compiler).__dict__) + + def test_config_hash_uses_native_contract_and_keeps_mutable_copies_uncached(self): + for release in (True, False): + config = SimpleNamespace(value=1, compute_hash=lambda: "release") + native = SimpleNamespace( + hash_config=lambda c: str(c.value), + hash_config_cached=MagicMock(return_value="cached"), + ) + + def import_config(name, *args, selected=native, **kwargs): + if name == "vllm.config": + return SimpleNamespace(utils=selected) + return builtins.__import__(name, *args, **kwargs) + + ns = load_nodes( + "vllm_ascend/utils.py", + {"get_config_hash"}, + { + "vllm_version_is": lambda _, selected_release=release: selected_release, + "config_utils": native, + "__builtins__": { + **vars(builtins), + "__import__": import_config, + }, + }, + ) + get_hash = ns["get_config_hash"] + self.assertEqual(get_hash(config), "release" if release else "1") + config.value = 2 + self.assertEqual(get_hash(config), "release" if release else "2") + self.assertEqual(get_hash(config, cached=True), "release" if release else "cached") + self.assertEqual(native.hash_config_cached.call_count, 0 if release else 1) + + def test_scheduler_payload_metrics_empty_prefill_decode_and_encoder(self): + @dataclass + class Details: + num_prefill_requests: int = 0 + num_prefill_tokens: int = 0 + num_decode_requests: int = 0 + num_decode_tokens: int = 0 + num_encoder_inputs: int = 0 + num_encoder_output_tokens: int = 0 + + for release in (True, False): + for logging, profiling in ((False, False), (True, False), (False, True)): + for scheduled in ({}, {"prefill": 6, "decode": 2}): + with self.subTest(release=release, logging=logging, profiling=profiling, scheduled=scheduled): + running, new_data = ["decode"], ["prefill"] + resumed: list[str] = [] + blocks: dict[str, Any] = {} + calls: list[tuple[Any, ...]] = [] + + def legacy(run, resume, tokens, spec, new_blocks, call_log=calls): + call_log.append((run, resume, tokens, spec, new_blocks)) + return "cached" + + def main(run, tokens, spec, new_blocks, iteration_details=None, call_log=calls): + call_log.append((run, tokens, spec, new_blocks)) + if iteration_details is not None and "decode" in tokens: + iteration_details.num_decode_requests = 1 + iteration_details.num_decode_tokens = tokens["decode"] + return "cached" + + updates = object() + scheduler = SimpleNamespace( + log_stats=logging, + observability_config=SimpleNamespace(enable_logging_iteration_details=True), + vllm_config=SimpleNamespace( + profiler_config=SimpleNamespace(profiler="torch" if profiling else None) + ), + max_num_encoder_input_tokens=32, + _make_cached_request_data=legacy if release else main, + _get_encoder_input_updates=MagicMock(return_value=updates), + ) + ns = load_nodes( + "vllm_ascend/core/scheduler_output.py", + {"make_cached_request_data"}, + { + "vllm_version_is": lambda _, selected_release=release: selected_release, + "IterationDetails": Details, + }, + ) + inputs = {"prefill": [0, 1]} if scheduled else {} + cached, metadata = ns["make_cached_request_data"]( + scheduler, + running, + resumed, + scheduled, + {}, + blocks, + new_data, + inputs, + 22 if scheduled else 32, + ) + self.assertEqual(cached, "cached") + self.assertIs(calls[0][0], running) + if release: + self.assertEqual(metadata, {}) + self.assertIs(calls[0][1], resumed) + scheduler._get_encoder_input_updates.assert_not_called() + else: + self.assertIs(metadata["encoder_input_updates"], updates) + scheduler._get_encoder_input_updates.assert_called_once_with(scheduled, new_data, inputs) + details = metadata["iteration_details"] + if logging or profiling: + self.assertEqual(details, Details(1, 6, 1, 2, 2, 10) if scheduled else Details()) + else: + self.assertIsNone(details) + + def test_main_engram_retains_slotless_launch_and_empty_input(self): + class Parent: + pass + + class Buffer: + shape = (0,) + + def new_empty(self, shape, **kwargs): + return shape + + ns = load_nodes( + "vllm_ascend/models/deepseek_v41/engram/hash_state.py", + {"AscendNgramHashState"}, + { + "NgramHashState": Parent, + "vllm_version_is": lambda _: False, + "torch": SimpleNamespace(int32="int32"), + }, + ) + state = ns["AscendNgramHashState"]() + state.swa_cache_module = SimpleNamespace(block_size=128) + state.multipliers = SimpleNamespace(shape=(3, 4)) + state.primes = SimpleNamespace(shape=(3, 2)) + self.assertFalse(state.use_slot_cache) + self.assertEqual(state.block_size, 128) + self.assertEqual(state.forward(Buffer(), None, None, None, None, None), (0, 3, 6)) + + def test_capture_guard_only_calls_release_monitor(self): + for release in (True, False): + monitor = SimpleNamespace(validate_cudagraph_capturing_enabled=MagicMock()) + ns = load_nodes( + "vllm_ascend/compilation/acl_graph.py", + {"validate_cudagraph_capturing_enabled"}, + { + "vllm_version_is": lambda _, selected_release=release: selected_release, + "monitor": monitor, + }, + ) + ns["validate_cudagraph_capturing_enabled"]() + self.assertEqual(monitor.validate_cudagraph_capturing_enabled.call_count, int(release)) + + def test_idle_dp_uses_full_speculative_query_length_on_main(self): + tree = ast.parse((ROOT / "vllm_ascend/worker/worker.py").read_text(encoding="utf-8")) + cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "NPUWorker") + method = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "execute_dummy_batch") + module = ast.Module(body=[method], type_ignores=[]) + for release in (True, False): + ns: dict[str, Any] = { + "vllm_version_is": lambda _, selected_release=release: selected_release, + "set_current_vllm_config": lambda _: nullcontext(), + "runtime_guard_idle_step": lambda fn: fn, + } + exec(compile(ast.fix_missing_locations(module), "", "exec"), ns) + worker = SimpleNamespace( + log_memory_stats=MagicMock(), + vllm_config=object(), + model_runner=SimpleNamespace( + decode_query_len=5, + _dummy_run=MagicMock(), + ), + ) + ns["execute_dummy_batch"](worker) + worker.model_runner._dummy_run.assert_called_once_with( + 1 if release else 5, uniform_decode=True, skip_gdn_state_update=True + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/ut/test_mrv2_utils.py b/tests/ut/test_mrv2_utils.py index 36ab60ef5a0d..28db64ef303a 100644 --- a/tests/ut/test_mrv2_utils.py +++ b/tests/ut/test_mrv2_utils.py @@ -56,7 +56,7 @@ def test_v1_request_is_rejected_after_upstream_drops_the_entry(monkeypatch): from vllm import envs # Upstream removes the env entry together with Model Runner V1. - monkeypatch.setattr(envs, "VLLM_USE_V2_MODEL_RUNNER", None) + monkeypatch.setattr(envs, "VLLM_USE_V2_MODEL_RUNNER", None, raising=False) monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") with pytest.raises(ValueError, match="VLLM_USE_V2_MODEL_RUNNER=0"): diff --git a/tests/ut/test_native_hardware_followup_contracts.py b/tests/ut/test_native_hardware_followup_contracts.py new file mode 100644 index 000000000000..a216865556e0 --- /dev/null +++ b/tests/ut/test_native_hardware_followup_contracts.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ast +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock, patch + +import pytest +import torch +import vllm + +from vllm_ascend.attention.sfa_v1 import AscendSFAImpl, AscendSFAMetadata, AscendSFAMetadataBuilder +from vllm_ascend.ops.fused_moe.routed_experts import EplbExpertTensorList +from vllm_ascend.utils import vllm_version_is + + +def _native_tree(path): + return ast.parse((Path(vllm.__file__).parent / path).read_text(encoding="utf-8")) + + +def _execute(nodes, scope): + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + tree = ast.Module(body=[future, *nodes], type_ignores=[]) + exec(compile(ast.fix_missing_locations(tree), "", "exec"), scope) + + +@pytest.mark.parametrize("with_draft", [False, True]) +@pytest.mark.parametrize("engine_fails", [False, True]) +def test_weight_transfer_invokes_native_replacement_callback(with_draft, engine_fails): + events = [] + release = vllm_version_is("0.33.0") + native = next( + node + for node in _native_tree("v1/worker/gpu/model_runner.py").body + if isinstance(node, ast.ClassDef) and node.name == "GPUModelRunner" + ) + native.body = [ + node + for node in native.body + if isinstance(node, ast.FunctionDef) + and node.name in ("reset_lora_state", "on_target_weights_replaced", "get_draft_model") + ] + assert any( + isinstance(node, ast.FunctionDef) + and node.name == ("reset_lora_state" if release else "on_target_weights_replaced") + for node in native.body + ) + base_methods = {"reset_lora_state": lambda self: events.append("reset release LoRA")} if release else {} + native.bases = [ast.Name(id="Base", ctx=ast.Load())] + native.decorator_list = [] + + class DraftModelSpeculator: + model = object() + + scope: dict[str, Any] = { + "Base": type("Base", (), base_methods), + "DraftModelSpeculator": DraftModelSpeculator, + "set_current_vllm_config": lambda config: nullcontext(), + "refresh_quantized_lm_heads": lambda model: events.append("refresh draft lm heads"), + "vllm_version_is": vllm_version_is, + } + _execute([native], scope) + runner = scope["GPUModelRunner"]() + runner.speculator = DraftModelSpeculator() if with_draft else None + runner.vllm_config = object() + worker_path = Path(__file__).resolve().parents[2] / "vllm_ascend/worker/worker.py" + worker_cls = next( + node + for node in ast.parse(worker_path.read_text(encoding="utf-8")).body + if isinstance(node, ast.ClassDef) and node.name == "NPUWorker" + ) + finish = next( + node for node in worker_cls.body if isinstance(node, ast.FunctionDef) and node.name == "finish_weight_update" + ) + _execute([finish], scope) + + def finish_base(): + events.append("finish base weights") + if engine_fails: + raise RuntimeError("weight engine failed") + + worker = SimpleNamespace( + _weight_update_active=True, + _check_weight_transfer_engine=lambda: None, + weight_transfer_engine=SimpleNamespace(finish_weight_update=finish_base), + model_runner=runner, + ) + if engine_fails: + with pytest.raises(RuntimeError, match="weight engine failed"): + scope["finish_weight_update"](worker) + assert events == ["finish base weights"] + assert worker._weight_update_active is True + else: + scope["finish_weight_update"](worker) + expected = ["finish base weights"] + if release: + expected.append("reset release LoRA") + if with_draft: + expected.append("refresh draft lm heads") + assert events == expected + assert worker._weight_update_active is False + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_native_eplb_copy_back_preserves_independent_experts(dtype): + dst = [EplbExpertTensorList(torch.zeros(2, 4, dtype=dtype) for _ in range(3)) for _ in range(2)] + src = [ + EplbExpertTensorList(torch.full((2, 4), float(i + offset + 1), dtype=dtype) for i in range(3)) + for offset in (0, 5) + ] + pointers = [[tensor.data_ptr() for tensor in weights] for weights in dst] + source_snapshot = [[tensor.clone() for tensor in weights] for weights in src] + assert dst[0].is_cuda is False + tree = _native_tree("distributed/eplb/rebalance_execute.py") + scope: dict[str, Any] = {} + if vllm_version_is("0.33.0"): + loop = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.For) + and ast.unparse(node.target) == "dst" + and "w[dst].copy_(b[dst]" in ast.unparse(node) + ) + scope.update(dest_indices=[2, 0], expert_weights=dst, expert_weights_buffers=src) + _execute([loop], scope) + pairs = [(2, 2), (0, 0)] + else: + copy = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "copy_rows") + _execute([copy], scope) + scope["copy_rows"](dst, [2, 0], src, [0, 2]) + scope["copy_rows"](dst, [0], dst, [0]) # Native add_model warmup is a no-op copy. + pairs = [(2, 0), (0, 2)] + for weights, source, snapshot, original_ptrs in zip(dst, src, source_snapshot, pointers): + assert [tensor.data_ptr() for tensor in weights] == original_ptrs + for dst_row, src_row in pairs: + torch.testing.assert_close(weights[dst_row], snapshot[src_row], rtol=0, atol=0) + assert torch.count_nonzero(weights[1]).item() == 0 + for tensor, original in zip(source, snapshot): + torch.testing.assert_close(tensor, original, rtol=0, atol=0) + + +def test_sfa_satisfies_native_sparse_guard_without_dense_prefill(): + release = vllm_version_is("0.33.0") + assert AscendSFAImpl.is_sparse is (not release) + assert AscendSFAImpl.supports_dense_mha_prefill is release + native = next( + node + for node in _native_tree("model_executor/layers/attention/mla_attention.py").body + if isinstance(node, ast.ClassDef) and node.name == "MLAAttention" + ) + init = next(node for node in native.body if isinstance(node, ast.FunctionDef) and node.name == "__init__") + guard = next( + ( + node + for node in init.body + if isinstance(node, ast.If) and ast.unparse(node.test) == "use_sparse and (not self.impl.is_sparse)" + ), + None, + ) + assert (guard is None) == release + if release: + return + state = SimpleNamespace( + impl=SimpleNamespace( + is_sparse=AscendSFAImpl.is_sparse, supports_dense_mha_prefill=AscendSFAImpl.supports_dense_mha_prefill + ), + attn_backend=SimpleNamespace(get_name=lambda: "ASCEND_SFA"), + supports_dense_mha_prefill=True, + ) + scope = dict(self=state, use_sparse=True, prefix="model.layers.0.self_attn.attn", logger=Mock()) + _execute([guard], scope) + prefill = next( + node + for node in init.body + if isinstance(node, ast.If) and ast.unparse(node.test).startswith("self.impl.is_sparse and") + ) + _execute([prefill], scope) + assert state.prefill_backend is None + state.impl.is_sparse = False + with pytest.raises(RuntimeError, match="requires a sparse attention backend"): + _execute([guard], scope) + + +@pytest.mark.parametrize("custom_metadata", [False, True]) +def test_sfa_metadata_constructor_keeps_sparse_backend_and_release_clone(custom_metadata): + release = vllm_version_is("0.33.0") + backend = Mock() if release else None + cloned_backend = object() + if backend is not None: + backend.clone.return_value = cloned_backend + layer = SimpleNamespace( + prefill_backend=backend, + q_lora_rank=2, + kv_lora_rank=2, + qk_nope_head_dim=2, + qk_rope_head_dim=2, + v_head_dim=2, + ) + config: Any = SimpleNamespace( + model_config=SimpleNamespace(dtype=torch.float32, hf_text_config=SimpleNamespace(mla_use_nope=False)), + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, decode_context_parallel_size=1, cp_kv_cache_interleave_size=1 + ), + compilation_config=SimpleNamespace(static_forward_context={"layer": layer}), + cache_config=SimpleNamespace(cache_dtype="auto"), + speculative_config=None, + ) + spec: Any = SimpleNamespace(block_size=128) + metadata_cls = type("CustomSFAMetadata", (AscendSFAMetadata,), {}) if custom_metadata else None + native_module = "vllm.model_executor.layers.attention.mla_attention" + with ( + patch(f"{native_module}.get_num_attention_heads_from_layers", return_value=1), + patch(f"{native_module}.get_dcp_group", return_value=SimpleNamespace(world_size=1)), + patch.object(AscendSFAMetadataBuilder, "determine_prefill_query_data_type", return_value=torch.float32), + patch.object(AscendSFAMetadataBuilder, "determine_chunked_prefill_workspace_size", return_value=4), + patch("vllm_ascend.attention.sfa_v1.is_pcp_decode_sharding_enabled", return_value=False), + ): + builder = AscendSFAMetadataBuilder(spec, ["layer"], config, torch.device("cpu"), metadata_cls) + assert layer.prefill_backend is backend + assert builder.metadata_cls is (metadata_cls if custom_metadata else AscendSFAMetadata) + assert builder.model_config is config.model_config + assert builder.kv_cache_spec is spec + assert builder.vllm_config is config + assert builder.device == torch.device("cpu") + assert builder.decode_threshold == builder.reorder_batch_threshold == 1 + if release: + assert backend is not None + backend.clone.assert_called_once_with() + assert builder._prefill_backend is cloned_backend + else: + assert not hasattr(builder, "_prefill_backend") + assert not hasattr(builder, "chunked_prefill_workspace") diff --git a/tests/ut/test_offline_dsd_tuner_config.py b/tests/ut/test_offline_dsd_tuner_config.py new file mode 100644 index 000000000000..3ccdb09ec812 --- /dev/null +++ b/tests/ut/test_offline_dsd_tuner_config.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ast +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock + +import pytest +from vllm.config import SpeculativeConfig + +from vllm_ascend.utils import vllm_version_is + + +class CapturedLLMConfig(BaseException): + def __init__(self, kwargs: dict[str, Any]): + self.kwargs = kwargs + + +@pytest.mark.parametrize("method", ["eagle", "eagle3"]) +@pytest.mark.parametrize("disable_padding", [False, True]) +def test_dsd_tuner_padding_matches_native_config(method, disable_padding): + """Exercise the example's LLM boundary without model loading or NPU work.""" + path = Path(__file__).resolve().parents[2] / "examples/offline_DSD_k_tuner/generate_config.py" + tree = ast.parse(path.read_text(encoding="utf-8")) + nodes = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name in ("get_acceptance_rate_per_pos", "build_compilation_config") + ] + assert len(nodes) == 2 + tokenizer = Mock() + + def capture_llm(**kwargs): + raise CapturedLLMConfig(kwargs) + + llm = Mock(side_effect=capture_llm) + scope: dict[str, Any] = { + "AutoTokenizer": tokenizer, + "get_samples": lambda *args: [], + "LLM": llm, + "vllm_version_is": vllm_version_is, + } + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + module = ast.Module(body=[future, *nodes], type_ignores=[]) + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), scope) + args = SimpleNamespace( + method=method, + disable_padded_drafter_batch=disable_padding, + model_dir="target", + enable_multimodal_chat=False, + eagle_dir="draft", + num_spec_tokens=3, + parallel_drafting=True, + tp=1, + enable_chunked_prefill=True, + enforce_eager=False, + gpu_memory_utilization=0.8, + max_model_len=1024, + max_vllm_batch_size=4, + allowed_local_media_path="", + cudagraph_capture_sizes=[1, 4], + ) + release = vllm_version_is("0.33.0") + field = "disable_padded_drafter_batch" + assert (field in SpeculativeConfig.__annotations__) == release + if disable_padding and not release: + with pytest.raises(ValueError, match="supported only with vLLM v0.33.0"): + scope["get_acceptance_rate_per_pos"](args) + tokenizer.from_pretrained.assert_not_called() + llm.assert_not_called() + else: + with pytest.raises(CapturedLLMConfig) as captured: + scope["get_acceptance_rate_per_pos"](args) + expected = { + "method": method, + "model": "draft", + "num_speculative_tokens": 3, + "parallel_drafting": True, + } + if release: + expected[field] = disable_padding + assert captured.value.kwargs["speculative_config"] == expected + assert captured.value.kwargs["compilation_config"] == { + "cudagraph_mode": "FULL_DECODE_ONLY", + "cudagraph_capture_sizes": [1, 4], + } diff --git a/tests/ut/test_platform.py b/tests/ut/test_platform.py index 6d15d2396b3a..38c6192b54ff 100644 --- a/tests/ut/test_platform.py +++ b/tests/ut/test_platform.py @@ -1,4 +1,5 @@ import importlib +from functools import partial from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -10,9 +11,9 @@ from vllm.platforms.interface import Platform from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.selector import AttentionSelectorConfig # type: ignore -from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher from tests.ut.base import TestBase +from tests.ut.graph_dispatch_fixture import CudagraphDispatcher from vllm_ascend.ascend_config import EplbConfig from vllm_ascend.ascend_forward_context import MoECommType, override_mrv2_in_profile_run from vllm_ascend.device.hardware_profile import get_hardware_profile @@ -31,6 +32,7 @@ COMPRESSED_TENSORS_METHOD, AscendDeviceType, dsv4_skips_indexer_topk, + vllm_version_is, ) @@ -392,6 +394,11 @@ def mock_vllm_config(): mock_vllm_config = MagicMock() mock_vllm_config.compilation_config = MagicMock() mock_vllm_config.compilation_config.reduced_cg_cap = None + mock_vllm_config.compilation_config.prefill_cudagraph_capture_token_sizes = None + mock_vllm_config.compilation_config.max_prefill_cudagraph_capture_tokens = None + mock_vllm_config.compilation_config.decode_cudagraph_capture_request_sizes = None + mock_vllm_config.compilation_config.max_decode_cudagraph_capture_requests = None + mock_vllm_config.compilation_config.max_decode_cudagraph_capture_tokens = None mock_vllm_config.model_config = MagicMock() mock_vllm_config.model_config.is_hybrid = False mock_vllm_config.model_config.is_encoder_decoder = False @@ -418,6 +425,7 @@ def mock_vllm_config(): ) mock_vllm_config.cache_config = MagicMock() mock_vllm_config.scheduler_config = MagicMock() + mock_vllm_config.scheduler_config.max_async_kv_loads_per_step = None mock_vllm_config.scheduler_config.max_num_seqs = None mock_vllm_config.scheduler_config.policy = "fcfs" mock_vllm_config.scheduler_config.async_scheduling = False @@ -920,6 +928,115 @@ def test_apply_config_platform_defaults_respects_explicit_sizes(self): self.assertIsNone(vllm_config.compilation_config.max_cudagraph_capture_size) self.assertEqual(vllm_config.compilation_config.cudagraph_capture_sizes, [1, 2, 4]) + @patch("vllm_ascend.platform.vllm_version_is", return_value=False) + def test_apply_config_platform_defaults_respects_explicit_split_sizes(self, _mock_version): + for name, value in ( + ("prefill_cudagraph_capture_token_sizes", [1, 3, 7]), + ("max_prefill_cudagraph_capture_tokens", 7), + ("decode_cudagraph_capture_request_sizes", [1, 3, 7]), + ("max_decode_cudagraph_capture_requests", 7), + ("max_decode_cudagraph_capture_tokens", 7), + ): + with self.subTest(name=name): + config = TestNPUPlatform.mock_vllm_config() + config.scheduler_config.max_num_seqs = 7 + config.compilation_config.max_cudagraph_capture_size = None + config.compilation_config.cudagraph_capture_sizes = None + setattr(config.compilation_config, name, value) + self.platform.apply_config_platform_defaults(config) + self.assertIsNone(config.compilation_config.max_cudagraph_capture_size) + self.assertIsNone(config.compilation_config.reduced_cg_cap) + self.assertEqual(getattr(config.compilation_config, name), value) + + def test_setup_compile_backend_restores_off_stride_split_sizes(self): + for max_num_seqs in (7, 77): + for mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_DECODE_ONLY): + with self.subTest(max_num_seqs=max_num_seqs, mode=mode): + config = TestNPUPlatform.mock_vllm_config() + config.compilation_config = CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, cudagraph_mode=mode + ) + config.scheduler_config.max_num_seqs = max_num_seqs + config.scheduler_config.max_num_batched_tokens = 1024 + config.uniform_decode_query_len = 1 + config.use_v2_model_runner = True + config.performance_mode = "balanced" + config.model_config.enforce_eager = False + config.additional_config = {"ascend_compilation_config": {}} + config.get_decode_cudagraph_query_lens = lambda q, count=max_num_seqs: {q: count} + config.get_draft_decode_cudagraph_request_sizes = lambda: [] + if not vllm_version_is("0.33.0"): + config._split_cudagraph_capture_sizes = partial( + VllmConfig._split_cudagraph_capture_sizes, config + ) + config._set_cudagraph_sizes = partial(VllmConfig._set_cudagraph_sizes, config) + self.platform.apply_config_platform_defaults(config) + config._set_cudagraph_sizes() + with patch("vllm_ascend.platform.enable_sp", return_value=False): + _setup_compile_backend(config, compile_backend="test_backend") + self.assertIn(max_num_seqs, config.compilation_config.cudagraph_capture_sizes) + self.assertEqual(config.compilation_config.max_cudagraph_capture_size, max_num_seqs) + config._set_cudagraph_sizes() + self.assertIn(max_num_seqs, config.compilation_config.cudagraph_capture_sizes) + + def test_setup_compile_backend_keeps_split_decode_grid_after_tp_alignment(self): + cases = ((1, 1, [4]), (4, 2, [16, 32]), (4, 1, [1, 2, 4, 8])) + for max_num_seqs, query_len, sizes in cases: + for mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_DECODE_ONLY): + for sequence_parallel in (False, True): + with self.subTest(seqs=max_num_seqs, query=query_len, mode=mode, sp=sequence_parallel): + config = TestNPUPlatform.mock_vllm_config() + config.compilation_config = CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + cudagraph_mode=mode, + cudagraph_capture_sizes=list(sizes), + ) + config.scheduler_config.max_num_seqs = max_num_seqs + config.scheduler_config.max_num_batched_tokens = 1024 + config.parallel_config.tensor_parallel_size = 4 + config.uniform_decode_query_len = query_len + config.use_v2_model_runner = True + config.performance_mode = "balanced" + config.speculative_config = None + config.model_config.enforce_eager = False + config.additional_config = {"ascend_compilation_config": {}} + config.get_decode_cudagraph_query_lens = lambda q, count=max_num_seqs: {q: count} + config.get_draft_decode_cudagraph_request_sizes = lambda: [] + if not vllm_version_is("0.33.0"): + config._split_cudagraph_capture_sizes = partial( + VllmConfig._split_cudagraph_capture_sizes, config + ) + config._set_cudagraph_sizes = partial(VllmConfig._set_cudagraph_sizes, config) + config._set_cudagraph_sizes() + decode_sizes = None + if not vllm_version_is("0.33.0"): + decode_sizes = config.compilation_config.decode_cudagraph_capture_request_sizes + if decode_sizes is not None: + decode_sizes = list(decode_sizes) + with patch("vllm_ascend.platform.enable_sp", return_value=sequence_parallel): + _setup_compile_backend( + config, compile_backend="test_backend", enable_shared_expert_dp=not sequence_parallel + ) + # Engine handshakes re-enter the native config resolver. + config._set_cudagraph_sizes() + config._set_cudagraph_sizes() + if not vllm_version_is("0.33.0"): + self.assertEqual( + config.compilation_config.decode_cudagraph_capture_request_sizes, decode_sizes + ) + prefill_sizes = config.compilation_config.prefill_cudagraph_capture_token_sizes + self.assertTrue(all(size % 4 == 0 for size in prefill_sizes)) + if max_num_seqs == 1: + self.assertEqual(prefill_sizes, [4]) + self.assertEqual(decode_sizes, [1]) + if mode == CUDAGraphMode.FULL_DECODE_ONLY: + self.assertEqual(config.compilation_config.cudagraph_capture_sizes, [1]) + else: + self.assertTrue( + all(size % 4 == 0 for size in config.compilation_config.cudagraph_capture_sizes) + ) + self.assertEqual(config.compilation_config.cudagraph_mode, mode) + def test_apply_config_platform_defaults_sets_sp_threshold(self): for enabled_option in ("enable_sp", "fuse_gemm_comms"): with self.subTest(enabled_option=enabled_option): @@ -1290,6 +1407,7 @@ def test_check_and_update_config_basic_config_update( vllm_config.parallel_config.tensor_parallel_size = 1 mock_init_recompute.return_value = MagicMock() vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None # Use importlib.reload to reload the platform module, ensuring the mocked init_ascend_config method is used. # Without this reload, when calling self.platform.check_and_update_config, @@ -1320,6 +1438,7 @@ def test_check_and_update_config_no_model_config_warning( vllm_config.parallel_config.tensor_parallel_size = 1 mock_init_recompute.return_value = MagicMock() vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None from vllm_ascend import platform @@ -1352,6 +1471,7 @@ def test_check_and_update_config_enforce_eager_mode( vllm_config.parallel_config.tensor_parallel_size = 1 mock_init_recompute.return_value = MagicMock() vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None from vllm_ascend import platform @@ -1397,6 +1517,7 @@ def test_check_and_update_config_unsupported_compilation_level( vllm_config.parallel_config.tensor_parallel_size = 1 mock_init_recompute.return_value = MagicMock() vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None vllm_config.compilation_config.mode = CompilationMode.DYNAMO_TRACE_ONCE @@ -1472,6 +1593,7 @@ def test_check_and_update_config_cache_config_block_size( vllm_config.parallel_config.tensor_parallel_size = 1 mock_init_recompute.return_value = MagicMock() vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None from vllm_ascend import platform @@ -1501,6 +1623,7 @@ def test_check_and_update_config_recompute_scheduler_rejects_pd_mixed_no_kv_tran vllm_config.parallel_config.prefill_context_parallel_size = 1 vllm_config.parallel_config.tensor_parallel_size = 1 vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None mock_init_recompute.return_value = MagicMock() from vllm_ascend import platform @@ -1536,6 +1659,7 @@ def test_check_and_update_config_recompute_scheduler_rejects_pd_mixed_kv_both( vllm_config.parallel_config.prefill_context_parallel_size = 1 vllm_config.parallel_config.tensor_parallel_size = 1 vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None mock_init_recompute.return_value = MagicMock() from vllm_ascend import platform @@ -1623,6 +1747,7 @@ def test_check_and_update_config_recompute_scheduler_accepts_kv_consumer( vllm_config.parallel_config.cp_kv_cache_interleave_size = 1 vllm_config.cache_config.block_size = 1 vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None vllm_config.scheduler_config.async_scheduling = False from vllm_ascend import platform @@ -2009,6 +2134,7 @@ def test_check_and_update_config_selects_dyntra_lb_scheduler( vllm_config.parallel_config.cp_kv_cache_interleave_size = 1 vllm_config.cache_config.block_size = 1 vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None vllm_config.scheduler_config.async_scheduling = False from vllm_ascend import platform @@ -2070,6 +2196,7 @@ def test_check_and_update_config_rejects_dyntra_lb_outside_disaggregated_d_node( vllm_config.parallel_config.cp_kv_cache_interleave_size = 1 vllm_config.cache_config.block_size = 1 vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None from vllm_ascend import platform @@ -2103,6 +2230,7 @@ def test_check_and_update_config_rejects_multinode_dyntra_lb_decoder( vllm_config.parallel_config.cp_kv_cache_interleave_size = 1 vllm_config.cache_config.block_size = 1 vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None from vllm_ascend import platform @@ -2137,6 +2265,7 @@ def test_check_and_update_config_balance_scheduler_rejects_pd_disaggregated_kv_p vllm_config.parallel_config.prefill_context_parallel_size = 1 vllm_config.parallel_config.tensor_parallel_size = 1 vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None mock_init_recompute.return_value = MagicMock() from vllm_ascend import platform @@ -2172,6 +2301,7 @@ def test_check_and_update_config_balance_scheduler_rejects_pd_disaggregated_kv_c vllm_config.parallel_config.prefill_context_parallel_size = 1 vllm_config.parallel_config.tensor_parallel_size = 1 vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None mock_init_recompute.return_value = MagicMock() from vllm_ascend import platform @@ -2490,6 +2620,7 @@ def test_check_and_update_config_v1_worker_class_selection( vllm_config.parallel_config.tensor_parallel_size = 1 mock_init_recompute.return_value = MagicMock() vllm_config.scheduler_config = MagicMock() + vllm_config.scheduler_config.max_async_kv_loads_per_step = None from vllm_ascend import platform @@ -2530,6 +2661,8 @@ def test_check_and_update_config_310p_no_custom_ops( mock_init_recompute.return_value = MagicMock() vllm_config.scheduler_config = MagicMock() + + vllm_config.scheduler_config.max_async_kv_loads_per_step = None from vllm_ascend import platform importlib.reload(platform) diff --git a/tests/ut/worker/test_attn_utils_v2.py b/tests/ut/worker/test_attn_utils_v2.py index b33882fa3c5b..366ffcd152d5 100644 --- a/tests/ut/worker/test_attn_utils_v2.py +++ b/tests/ut/worker/test_attn_utils_v2.py @@ -1463,7 +1463,10 @@ def _ascend_allocate_kv_cache( _device: torch.device, _layout: Any, _kernel_block_sizes: list[int], + *, + allocate=None, ) -> dict[str, Any]: + assert allocate is None del _layout raw_tensors = attn_utils._allocate_kv_cache( _kv_cache_config, diff --git a/tests/ut/worker/test_idle_dummy_config.py b/tests/ut/worker/test_idle_dummy_config.py index bd650bcd390c..9cbcebd504f1 100644 --- a/tests/ut/worker/test_idle_dummy_config.py +++ b/tests/ut/worker/test_idle_dummy_config.py @@ -25,7 +25,7 @@ def dummy(num_tokens, **kwargs): worker = SimpleNamespace( vllm_config=config, log_memory_stats=Mock(), - model_runner=SimpleNamespace(uniform_decode_query_len=2, _dummy_run=dummy), + model_runner=SimpleNamespace(uniform_decode_query_len=2, decode_query_len=2, _dummy_run=dummy), ) with set_current_vllm_config(outer_config): if fails: diff --git a/tests/ut/worker/test_model_runner_v2.py b/tests/ut/worker/test_model_runner_v2.py index 8c6b67605cab..45421aa9ee12 100644 --- a/tests/ut/worker/test_model_runner_v2.py +++ b/tests/ut/worker/test_model_runner_v2.py @@ -27,6 +27,7 @@ from vllm_ascend.attention.attention_v1 import AscendAttentionBackend from vllm_ascend.attention.fa_v1 import AscendFlashAttentionBackend from vllm_ascend.attention.mla_v1 import AscendMLABackend, AscendMLAImpl +from vllm_ascend.utils import vllm_version_is from vllm_ascend.worker.v2.input_batch import AscendInputBatch, AscendInputBuffers from vllm_ascend.worker.v2.model_runner import NPUModelRunner from vllm_ascend.worker.v2.pcp_manager import AscendPCPManager @@ -1280,7 +1281,10 @@ def _prepare_kda(context, maximum): ) as create_kvpp, ): runner.initialize_kv_cache(kv_cache_config) - seen["factory"](runner.vllm_config, torch.device("cpu"), CUDAGraphMode.FULL, 1) + query_kwargs: dict[str, Any] = ( + {"decode_query_len": 1} if vllm_version_is("0.33.0") else {"decode_query_lens": {1: 8}} + ) + seen["factory"](runner.vllm_config, torch.device("cpu"), CUDAGraphMode.FULL, **query_kwargs) prepare_kda.assert_called_once_with(runner.compilation_config.static_forward_context, 8) assert runner.model_state.device_metadata == ( diff --git a/tests/ut/worker/test_model_runner_v2_kv_cache.py b/tests/ut/worker/test_model_runner_v2_kv_cache.py index eacfefa08110..8ecd951346f2 100644 --- a/tests/ut/worker/test_model_runner_v2_kv_cache.py +++ b/tests/ut/worker/test_model_runner_v2_kv_cache.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext from unittest.mock import MagicMock, patch import pytest @@ -15,7 +15,8 @@ @pytest.mark.parametrize("is_profiling", [False, True]) -def test_initialize_preserves_connector_containers_and_flattens_runner_cache(is_profiling): +@pytest.mark.parametrize("has_trackers", [False, True]) +def test_initialize_preserves_connector_containers_and_flattens_runner_cache(is_profiling, has_trackers): k, v, conv, ssm, single = [torch.empty(3, 4) for _ in range(5)] other = torch.empty(3, 4, device="meta") caches = {"attention": (k, None, v), "mamba": [conv, ssm], "single": single, "other": other} @@ -30,6 +31,9 @@ def test_initialize_preserves_connector_containers_and_flattens_runner_cache(is_ # the test exercises the KV-connector path only (main-only branch). runner.vllm_config.aux_output_config.enabled = False runner.jit_warmup_registry.activate.side_effect = nullcontext + stale_trackers = [torch.empty(3, dtype=torch.int32)] + runner.replayssm_block_copy_tensors = stale_trackers + trackers = [torch.zeros(3, dtype=torch.int32), torch.ones(3, dtype=torch.int32)] if has_trackers else [] config = KVCacheConfig(num_blocks=3, kv_cache_tensors=[], kv_cache_groups=[]) registration_order = [] runner._register_sparse_kv_caches.side_effect = lambda _: registration_order.append("offload") @@ -37,6 +41,11 @@ def test_initialize_preserves_connector_containers_and_flattens_runner_cache(is_ def register_connector(*_args): registration_order.append("connector") + tracker_context: AbstractContextManager[MagicMock | None] + if not vllm_version_is("0.33.0"): + tracker_context = patch.object(upstream, "get_replayssm_block_copy_tensors", return_value=trackers) + else: + tracker_context = nullcontext(None) with ( patch.object(upstream, "init_attn_backend", return_value=([], MagicMock(), [])), patch.object(upstream, "maybe_create_adaptive_verification_manager", return_value=None), @@ -48,11 +57,18 @@ def register_connector(*_args): patch.object(upstream, "ModelCudaGraphManager"), patch.object(upstream, "check_attention_cp_compatibility"), patch.object(upstream, "init_kv_cache", return_value=caches), + tracker_context as collect_trackers, patch.object(upstream, "get_kv_connector", side_effect=register_connector) as connector, ): upstream.GPUModelRunner.initialize_kv_cache(runner, config, is_profiling=is_profiling) assert [id(tensor) for tensor in runner.kv_caches] == [id(tensor) for tensor in (k, v, conv, ssm, single)] + if collect_trackers is not None: + collect_trackers.assert_called_once_with(runner.compilation_config.static_forward_context) + assert runner.replayssm_block_copy_tensors is trackers + assert type(runner.replayssm_block_copy_tensors) is list + else: + assert runner.replayssm_block_copy_tensors is stale_trackers runner._register_sparse_kv_caches.assert_called_once_with(caches) if is_profiling: assert registration_order == ["offload"] @@ -118,8 +134,14 @@ def test_mrv2_block_copy_preserves_segmented_mamba_storage(): conv = storage[:12].view(3, 4) ssm = storage[12:].view(3, 8) before_conv, before_ssm = conv.clone(), ssm.clone() - upstream.copy_kv_cache_blocks_inplace([conv, ssm, conv, ssm], 3, [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)]) + ring_start = torch.tensor([10, 20, 30], dtype=torch.int32) + num_committed = torch.tensor([1, 2, 3], dtype=torch.int32) + upstream.copy_kv_cache_blocks_inplace( + [conv, ssm, conv, ssm, ring_start, num_committed], 3, [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)] + ) torch.testing.assert_close(conv[2], before_conv[0]) torch.testing.assert_close(ssm[2], before_ssm[0]) torch.testing.assert_close(conv[:2], before_conv[:2]) torch.testing.assert_close(ssm[:2], before_ssm[:2]) + torch.testing.assert_close(ring_start, torch.tensor([10, 20, 10], dtype=torch.int32)) + torch.testing.assert_close(num_committed, torch.tensor([1, 2, 1], dtype=torch.int32)) diff --git a/tests/ut/worker/test_model_runner_v2_mamba.py b/tests/ut/worker/test_model_runner_v2_mamba.py index f5119b7140db..093c30506ef2 100644 --- a/tests/ut/worker/test_model_runner_v2_mamba.py +++ b/tests/ut/worker/test_model_runner_v2_mamba.py @@ -25,6 +25,7 @@ from vllm_ascend.core.kv_cache_interface import AscendMLAAttentionSpec from vllm_ascend.device.hardware import AscendDeviceType from vllm_ascend.device.hardware_profile import get_hardware_profile +from vllm_ascend.utils import vllm_version_is from vllm_ascend.worker.device_metadata import DeviceMetadataStage, DeviceMetadataTask, TargetDeviceMetadata from vllm_ascend.worker.v2 import attn_utils from vllm_ascend.worker.v2.attn_utils import ( @@ -303,16 +304,34 @@ def test_ensure_align_ctx_lazily_recreates_mamba_ctx(): patch("vllm_ascend.worker.v2.model_states.mamba_hybrid.validate_mamba_state_copy_funcs"), patch( "vllm_ascend.worker.v2.model_states.mamba_hybrid.MambaSpecDecodeGPUContext.create", + autospec=True, return_value=ctx, ) as mock_create, ): - result = state._ensure_align_ctx(kv_cache_config, [0], (torch.zeros(3, dtype=torch.int32),)) + block_table = torch.zeros(3, dtype=torch.int32) + result = state._ensure_align_ctx(kv_cache_config, [0], (block_table,)) + ctx.is_initialized = True + assert state._ensure_align_ctx(kv_cache_config, [0], (block_table,)) is ctx assert result is ctx assert state._mamba_ctx is ctx assert state._mamba_state_copy_funcs is copy_funcs mock_create.assert_called_once() - ctx.initialize_from_forward_context.assert_called_once() + # Enforce the installed upstream signature, rather than accepting any + # keyword/omitted argument through an unconstrained MagicMock. + if not vllm_version_is("0.33.0"): + assert ( + mock_create.call_args.kwargs["forward_context"] + is state.vllm_config.compilation_config.static_forward_context + ) + else: + assert "forward_context" not in mock_create.call_args.kwargs + ctx.initialize_from_forward_context.assert_called_once_with( + kv_cache_config, + state.vllm_config.compilation_config.static_forward_context, + copy_funcs, + [block_table], + ) @patch("vllm.v1.worker.mamba_utils.precopy_mamba_align_fused_kernel") diff --git a/tests/ut/worker/test_mtp_pcp_speculator_v2.py b/tests/ut/worker/test_mtp_pcp_speculator_v2.py index 60fb62641211..c5fa119e66c0 100644 --- a/tests/ut/worker/test_mtp_pcp_speculator_v2.py +++ b/tests/ut/worker/test_mtp_pcp_speculator_v2.py @@ -3,7 +3,7 @@ # Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. import json -from contextlib import nullcontext +from contextlib import ExitStack, nullcontext from copy import deepcopy from dataclasses import replace as dataclass_replace from types import SimpleNamespace @@ -21,6 +21,7 @@ VllmConfig, replace, ) +from vllm.config.compilation import CUDAGraphMode from vllm.platforms import current_platform from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp @@ -119,10 +120,12 @@ def eplb_worker_config(tmp_path): "vllm.model_executor.models.ModelRegistry.inspect_model_cls", return_value=(model_info, model_info.architecture), ), - patch("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", True), - patch("vllm.config.vllm.HAS_TRITON", True), + ExitStack() as legacy_runner, patch("vllm.config.parallel.current_platform.is_cuda_alike", return_value=True), ): + if vllm_version_is("0.33.0"): + legacy_runner.enter_context(patch("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", True)) + legacy_runner.enter_context(patch("vllm.config.vllm.HAS_TRITON", True)) # Keep the model dtype float-compatible with the bfloat16 KV cache the # draft loader override below sets: vLLM main rejects a mismatched # explicit float KV dtype. @@ -322,7 +325,7 @@ def _mla_pcp_decode_metadata(num_tokens, query_ends): def test_pcp_draft_mla_query_lengths_match_real_decode_rows(num_tokens, query_ends, expected): target = _mla_pcp_decode_metadata(num_tokens, query_ends) target_metadata = {"mtp": target} - draft = AscendMTPSpeculator._get_pcp_draft_prefill_metadata(target_metadata)["mtp"] + draft = AscendMTPSpeculator._get_draft_prefill_metadata(target_metadata)["mtp"] # This is the FIA TND contract that failed on PCP1 with Q=4, offsets=[4, 8]. assert draft.decode.actual_seq_lengths_q == expected @@ -344,10 +347,42 @@ def test_pcp_draft_metadata_preserves_other_attention_paths(): prefill_only.decode = None other_backend = SimpleNamespace(num_decode_tokens=4, decode=SimpleNamespace(actual_seq_lengths_q=[4, 8])) metadata = {"prefill": prefill_only, "other": other_backend, "missing": None} - assert AscendMTPSpeculator._get_pcp_draft_prefill_metadata(metadata) == metadata + assert AscendMTPSpeculator._get_draft_prefill_metadata(metadata) == metadata assert other_backend.decode.actual_seq_lengths_q == [4, 8] +@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]) +@pytest.mark.parametrize("query_ends", [[2, 4, 6, 8], [2]]) +def test_non_pcp_draft_prefill_normalizes_main_target_graph_padding(mode, query_ends): + speculator = object.__new__(AscendMTPSpeculator) + speculator.pcp_manager = None + speculator.draft_attn_layer_names = {"mtp"} + target = _mla_pcp_decode_metadata(2, query_ends) + other = SimpleNamespace() + target_metadata = {"mtp": target, "target_only": other} + slots = {"mtp": target.slot_mapping} + + with patch.object(speculator_module.TargetDependentARSpeculator, "_prefill") as parent: + speculator._prefill(1, 8, target_metadata, slots, None, mode) + + passed_metadata = parent.call_args.args[2] + draft = passed_metadata["mtp"] + expected = query_ends if vllm_version_is("0.33.0") or mode == CUDAGraphMode.FULL else [2] * len(query_ends) + assert draft.decode.actual_seq_lengths_q == expected + assert set(passed_metadata) == {"mtp"} + assert target.decode.actual_seq_lengths_q == query_ends + assert target_metadata == {"mtp": target, "target_only": other} + assert draft.slot_mapping is target.slot_mapping + assert draft.decode.seq_lens_list is target.decode.seq_lens_list + assert parent.call_args.args[0:2] == (1, 8) + assert parent.call_args.args[3] is slots + assert parent.call_args.args[5] == mode + if expected != query_ends: + assert draft is not target and draft.decode is not target.decode + else: + assert draft is target + + @pytest.mark.parametrize("pcp_enabled", [False, True]) @pytest.mark.parametrize("mismatch", ["tokens", "uniform"]) @pytest.mark.parametrize("dummy_run", [False, True]) diff --git a/tests/ut/worker/test_pcp_manager_v2.py b/tests/ut/worker/test_pcp_manager_v2.py index 48a5196650de..f67793cc7c91 100644 --- a/tests/ut/worker/test_pcp_manager_v2.py +++ b/tests/ut/worker/test_pcp_manager_v2.py @@ -155,6 +155,7 @@ def test_decode_sharding_does_not_add_a_parallel_config_field(): speculative_config=SimpleNamespace(method="mtp"), compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE), cache_config=None, + scheduler_config=SimpleNamespace(max_async_kv_loads_per_step=None), ) ascend_config = SimpleNamespace( scheduler_config=SimpleNamespace( diff --git a/tests/ut/worker/test_pcp_profile_budget.py b/tests/ut/worker/test_pcp_profile_budget.py index 621d926bfb4b..3595fc4f51ed 100644 --- a/tests/ut/worker/test_pcp_profile_budget.py +++ b/tests/ut/worker/test_pcp_profile_budget.py @@ -18,7 +18,9 @@ def make_runner(tokens=16384, pcp=16, requests=64, query=4, replicated=False): runner.max_num_reqs = requests runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp)) runner.speculative_config = SimpleNamespace( - method="dspark" if replicated else "mtp", num_speculative_tokens=query - 1 + method="dspark" if replicated else "mtp", + num_speculative_tokens=query - 1, + max_num_new_slots_for_drafting=0, ) runner.speculator = None # Non-last PP stages have no local drafter. return runner diff --git a/tests/ut/worker/test_v2_dummy_padding.py b/tests/ut/worker/test_v2_dummy_padding.py index 709a4f715e34..a58c0e1d7eae 100644 --- a/tests/ut/worker/test_v2_dummy_padding.py +++ b/tests/ut/worker/test_v2_dummy_padding.py @@ -49,12 +49,23 @@ def runner(buffers): def _make_graph_manager(buffers, mode, uniform, capture_size): manager = ModelAclGraphManager.__new__(ModelAclGraphManager) manager.vllm_config = SimpleNamespace( - speculative_config=None, cache_config=SimpleNamespace(use_kda_recoverssm=False) + speculative_config=None, + cache_config=SimpleNamespace(use_kda_recoverssm=False), + uniform_decode_query_len=uniform, ) - manager.compilation_config = SimpleNamespace(cudagraph_capture_sizes=[capture_size], max_cudagraph_capture_size=64) + manager.compilation_config = SimpleNamespace( + cudagraph_capture_sizes=[capture_size], + max_cudagraph_capture_size=64, + prefill_cudagraph_capture_token_sizes=[capture_size], + decode_cudagraph_capture_request_sizes=[(capture_size + uniform - 1) // uniform], + max_decode_cudagraph_capture_tokens=None, + ) + manager.decode_request_sizes = None manager.cudagraph_mode = mode manager.decode_query_len = uniform manager.max_num_reqs = len(buffers.seq_lens_cpu) + manager.decode_max_reqs = {uniform: manager.max_num_reqs} + manager._decode_descs = {} manager.varlen_decode = False manager.ubatch_runner = None manager.lora_capture_cases = [0] diff --git a/tests/ut/worker/test_worker_v1.py b/tests/ut/worker/test_worker_v1.py index 11ece7a8ce6f..aff88df73132 100644 --- a/tests/ut/worker/test_worker_v1.py +++ b/tests/ut/worker/test_worker_v1.py @@ -913,6 +913,20 @@ def test_init_device_selects_mm_encoder_runner_under_cuda_wrapper(self): self.assertIs(worker.model_runner, mm_encoder_runner.return_value) report_usage_stats.assert_not_called() + def test_memory_utilization_resolves_default_and_preserves_explicit_budget(self): + from vllm_ascend.worker.worker import NPUWorker + + worker = NPUWorker.__new__(NPUWorker) + worker.cache_config = CacheConfig(gpu_memory_utilization=0.5) + self.assertEqual(worker._get_gpu_memory_utilization(), 0.5) + worker.cache_config = CacheConfig() + expected = worker.cache_config.gpu_memory_utilization + if expected is None: + worker.cache_config.enable_extensible_kv_cache = False + expected = worker.cache_config.DEFAULT_GPU_MEMORY_UTILIZATION + self.assertEqual(worker._get_gpu_memory_utilization(), expected) + self.assertEqual(2000 * worker._get_gpu_memory_utilization(), 2000 * expected) + def test_profile_start_stop(self): """Test profile method start and stop""" from vllm_ascend.worker.worker import NPUWorker @@ -1135,6 +1149,7 @@ def test_execute_dummy_batch(self): worker.vllm_config = MagicMock() mock_model_runner = MagicMock() mock_uniform_decode_query_len = mock_model_runner.uniform_decode_query_len + mock_model_runner.decode_query_len = mock_uniform_decode_query_len worker.model_runner = mock_model_runner # Test execute_dummy_batch @@ -2389,7 +2404,8 @@ def test_finish_weight_update_resets_lora_state_after_base_weights(self): worker = self._make_worker(engine=engine) worker._weight_update_active = True - worker.finish_weight_update() + with patch("vllm_ascend.worker.worker.vllm_version_is", return_value=True): + worker.finish_weight_update() engine.finish_weight_update.assert_called_once_with() worker.model_runner.reset_lora_state.assert_called_once_with() diff --git a/tests/ut/worker/test_worker_v2.py b/tests/ut/worker/test_worker_v2.py index d1361e759da2..ed3a352e311b 100644 --- a/tests/ut/worker/test_worker_v2.py +++ b/tests/ut/worker/test_worker_v2.py @@ -14,6 +14,7 @@ from vllm.v1.kv_cache_layout import KVCacheLayout from tests.ut.base import TestBase +from vllm_ascend.utils import vllm_version_is @pytest.mark.parametrize("enforce_eager", [False, True]) @@ -304,7 +305,9 @@ def test_execute_model_pooling_uses_pool_output(self, mock_get_pp_group, mock_ge worker.profiler = None worker._pp_send_work = [] - pooling_output = ModelRunnerOutput(req_ids=[], req_id_to_index={}) + pooling_output = ModelRunnerOutput( + req_ids=[], **({"req_id_to_index": {}} if vllm_version_is("0.33.0") else {}) + ) worker.model_runner.execute_model.return_value = None worker.model_runner.pool.return_value = pooling_output diff --git a/tests/ut/worker/v2/test_aclgraph.py b/tests/ut/worker/v2/test_aclgraph.py index a0fe08d1bb9e..97fa1454fd35 100644 --- a/tests/ut/worker/v2/test_aclgraph.py +++ b/tests/ut/worker/v2/test_aclgraph.py @@ -9,12 +9,13 @@ import pytest import torch from vllm.config.compilation import CUDAGraphMode -from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( - SpeculatorCudaGraphManager, -) +from vllm_ascend.ascend_forward_context import MoECommType +from vllm_ascend.utils import vllm_version_is from vllm_ascend.worker.v2.spec_decode.target_dependent_ar.aclgraph import ( AutoRegressiveAclGraphManager, + CudaGraphManager, + SpeculatorCudaGraphManager, ) @@ -26,6 +27,7 @@ def _parent_init( decode_query_len, lora_capture_cases=None, varlen_decode=False, + **kwargs, ): """Mock the parent graph manager initialization.""" self._capture_descs = [object()] @@ -36,6 +38,12 @@ def _graph_config(pcp_size: int = 1) -> SimpleNamespace: return SimpleNamespace(parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size)) +def _query_shape(query_len): + if vllm_version_is("0.33.0"): + return {"decode_query_len": query_len} + return {"decode_query_lens": {query_len: 8}} + + def test_init_draft_prefill_graph(): """Verify initialization for the draft prefill graph.""" with ( @@ -51,7 +59,7 @@ def test_init_draft_prefill_graph(): patch("vllm_ascend.worker.v2.spec_decode.target_dependent_ar.aclgraph.set_draft_graph_params") as set_decode, ): manager = AutoRegressiveAclGraphManager( - _graph_config(), torch.device("cpu"), CUDAGraphMode.FULL, decode_query_len=2 + _graph_config(), torch.device("cpu"), CUDAGraphMode.FULL, **_query_shape(2) ) assert manager.capture_sizes == [4, 8] assert manager.is_draft_model_prefill is True @@ -79,7 +87,7 @@ def test_sharded_pcp_draft_prefill_avoids_full_graphs(decode_query_len, mode, ex ), ): manager = AutoRegressiveAclGraphManager( - _graph_config(pcp_size=2), torch.device("cpu"), mode, decode_query_len=decode_query_len + _graph_config(pcp_size=2), torch.device("cpu"), mode, **_query_shape(decode_query_len) ) assert manager.cudagraph_mode == expected @@ -98,12 +106,105 @@ def test_init_draft_decode_graph(): ) as set_prefill, patch("vllm_ascend.worker.v2.spec_decode.target_dependent_ar.aclgraph.set_draft_graph_params") as set_decode, ): - manager = AutoRegressiveAclGraphManager(object(), torch.device("cpu"), CUDAGraphMode.FULL, decode_query_len=1) + manager = AutoRegressiveAclGraphManager(object(), torch.device("cpu"), CUDAGraphMode.FULL, **_query_shape(1)) assert manager.is_draft_model_prefill is False set_decode.assert_called_once_with([4, 8]) set_prefill.assert_not_called() +@pytest.mark.skipif(vllm_version_is("0.33.0"), reason="Independent draft request grids require main's graph API") +@pytest.mark.parametrize( + "draft_moe,method,query_len,mode,limit,expected_max_requests", + [ + (True, MoECommType.MC2, 1, CUDAGraphMode.FULL_DECODE_ONLY, 16, 16), + (True, MoECommType.MC2, 1, CUDAGraphMode.FULL_DECODE_ONLY, 8, 8), + (True, MoECommType.FUSED_MC2, 1, CUDAGraphMode.FULL_DECODE_ONLY, 512, 512), + (True, MoECommType.ALLGATHER, 1, CUDAGraphMode.FULL_DECODE_ONLY, 512, 512), + (False, MoECommType.MC2, 1, CUDAGraphMode.FULL_DECODE_ONLY, 16, 512), + (True, MoECommType.MC2, 4, CUDAGraphMode.FULL_DECODE_ONLY, 16, 512), + (True, MoECommType.MC2, 1, CUDAGraphMode.NONE, 16, 512), + (True, MoECommType.ALLTOALL, 1, CUDAGraphMode.FULL_DECODE_ONLY, 512, 0), + ], +) +def test_draft_decode_capture_respects_moe_comm_policy( + draft_moe, method, query_len, mode, limit, expected_max_requests +): + sizes = [1, 4, 8, 16, 32, 512] + original_sizes = sizes.copy() + parent_kwargs: dict[str, Any] = {} + + def parent_init(self, *args, **kwargs): + parent_kwargs.update(kwargs) + _parent_init(self, *args, **kwargs) + + with ( + patch.object(SpeculatorCudaGraphManager, "__init__", parent_init), + patch.object(SpeculatorCudaGraphManager, "needs_capture", return_value=False), + patch( + "vllm_ascend.worker.v2.spec_decode.target_dependent_ar.aclgraph.collect_sorted_captured_token_sizes", + return_value=[], + ), + patch( + "vllm_ascend.worker.v2.spec_decode.target_dependent_ar.aclgraph.is_drafter_moe_model", + return_value=draft_moe, + ), + patch( + "vllm_ascend.worker.v2.spec_decode.target_dependent_ar.aclgraph.select_moe_comm_method", + side_effect=lambda size, _: method if size <= limit else MoECommType.ALLTOALL, + ), + ): + AutoRegressiveAclGraphManager( + _graph_config(), + torch.device("cpu"), + mode, + decode_query_lens={query_len: 512}, + decode_request_sizes=sizes, + ) + compatible_sizes = parent_kwargs["decode_request_sizes"] + assert max(compatible_sizes, default=0) == expected_max_requests + assert sizes == original_sizes + + # Exercise native descriptor creation and dispatch, including the eager + # boundary, without allocating a device graph or changing the config grid. + compilation = SimpleNamespace( + prefill_cudagraph_capture_token_sizes=[4, 8, 12, 16], + decode_cudagraph_capture_request_sizes=[1, 2, 3, 4], + max_decode_cudagraph_capture_tokens=None, + max_cudagraph_capture_size=4096, + ) + native: Any = SimpleNamespace( + compilation_config=compilation, + cudagraph_mode=mode, + max_num_reqs=512, + decode_query_len=query_len, + decode_max_reqs={query_len: 512}, + decode_request_sizes=compatible_sizes, + varlen_decode=False, + lora_capture_cases=[0], + vllm_config=SimpleNamespace(uniform_decode_query_len=4, cache_config=SimpleNamespace(use_kda_recoverssm=False)), + _capture_descs={}, + _decode_descs={}, + _candidates={}, + _graphs_captured=True, + _maybe_ubatch_twin=lambda _: None, + _resolve_effective_loras=lambda count: count, + ) + CudaGraphManager._init_candidates(native) + captured = [desc for descs in native._capture_descs.values() for desc in descs] + if mode == CUDAGraphMode.NONE: + assert not captured + else: + assert max((desc.num_reqs for desc in captured), default=0) == expected_max_requests + if draft_moe and query_len == 1 and method == MoECommType.MC2 and mode != CUDAGraphMode.NONE: + at_limit = CudaGraphManager.dispatch(native, limit, limit, uniform_token_count=1, num_active_loras=0) + above_limit = CudaGraphManager.dispatch(native, limit + 1, limit + 1, uniform_token_count=1, num_active_loras=0) + assert at_limit.cg_mode == CUDAGraphMode.FULL + assert above_limit.cg_mode == CUDAGraphMode.NONE + assert above_limit.num_reqs == limit + 1 + assert above_limit.num_tokens == limit + 1 + assert compilation.decode_cudagraph_capture_request_sizes == [1, 2, 3, 4] + + def test_run_fullgraph_uses_updatable_graph(): """Verify full-graph replay uses the updatable graph path.""" manager = AutoRegressiveAclGraphManager.__new__(AutoRegressiveAclGraphManager) diff --git a/tests/ut/worker/v2/test_entropy_verify.py b/tests/ut/worker/v2/test_entropy_verify.py index 66230d274dc0..925bf3ae231c 100644 --- a/tests/ut/worker/v2/test_entropy_verify.py +++ b/tests/ut/worker/v2/test_entropy_verify.py @@ -122,6 +122,19 @@ def test_invalid_entropy_parameters_fail_before_launch(inputs, launches, thresho mock.assert_not_called() +@pytest.mark.parametrize("skip_masked_vocab_blocks", [False, True]) +def test_masked_vocab_hint_preserves_npu_resample(monkeypatch, inputs, launches, skip_masked_vocab_blocks): + set_config(monkeypatch) + sampled, num_sampled = utils.rejection_sample(**inputs, skip_masked_vocab_blocks=skip_masked_vocab_blocks) + launches["_compute_block_stats_kernel"].assert_called_once() + launches["_probabilistic_rejection_kernel"].assert_called_once() + launches["resample"].assert_called_once() + args = launches["resample"].call_args.args + assert args[0] is sampled + assert args[1] is num_sampled + assert args[2] is inputs["target_logits"] + + def test_disabled_entropy_preserves_upstream_block(monkeypatch, inputs, launches): set_config(monkeypatch) utils.rejection_sample(**inputs, use_block_verification=True) diff --git a/tests/ut/worker/v2/test_lmhead_tp_utils.py b/tests/ut/worker/v2/test_lmhead_tp_utils.py index a96456949f97..a192bdd3fbe9 100644 --- a/tests/ut/worker/v2/test_lmhead_tp_utils.py +++ b/tests/ut/worker/v2/test_lmhead_tp_utils.py @@ -349,7 +349,7 @@ def propose(self, *args, **kwargs): draft.__dict__.update(make_drafter(mixin).__dict__) draft.dp_size = dp_size draft.pcp_manager = object() if pcp_enabled else None - draft._get_pcp_draft_prefill_metadata = lambda metadata: metadata + draft._get_draft_prefill_metadata = lambda metadata: metadata batch = SimpleNamespace(num_reqs=2, num_tokens_after_padding=24) state = None if shared_reqs is None else SimpleNamespace(num_reqs=shared_reqs) assert draft.propose(batch, *([None] * 10), dp_sync_state=state) == "drafted" diff --git a/tests/ut/worker/v2/test_spec_decode.py b/tests/ut/worker/v2/test_spec_decode.py index 031546ff01e9..6ffadcf987b4 100644 --- a/tests/ut/worker/v2/test_spec_decode.py +++ b/tests/ut/worker/v2/test_spec_decode.py @@ -1,14 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +import ast import sys +from pathlib import Path from types import ModuleType, SimpleNamespace +from typing import Any from unittest.mock import MagicMock import pytest import torch -from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode +from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, SpeculativeConfig +from vllm_ascend.worker.v2 import spec_decode as v2_spec_decode from vllm_ascend.worker.v2.spec_decode import init_speculator from vllm_ascend.worker.v2.spec_decode.rejection_sampler_utils import rejection_sample @@ -53,7 +57,10 @@ def _make_config( use_gemma4_mtp=MagicMock(return_value=use_gemma4_mtp), use_step3p5_mtp=MagicMock(return_value=use_step3p5_mtp), use_multi_module_mtp=MagicMock(return_value=False), - draft_model_config=SimpleNamespace(architectures=architectures or []), + draft_model_config=SimpleNamespace( + architectures=architectures or [], + hf_config=SimpleNamespace(model_type="step3p5_mtp" if use_step3p5_mtp else None), + ), ) return SimpleNamespace(speculative_config=speculative_config) @@ -288,3 +295,35 @@ def test_init_speculator_skips_special_mtp(): match="mtp is not supported yet", ): init_speculator(config, torch.device("cpu")) + + +@pytest.mark.parametrize("use_310p", [False, True]) +@pytest.mark.parametrize("model_type", ["qwen3_5_mtp", "step3p5_mtp", "gemma4_mtp"]) +def test_mtp_routing_uses_actual_native_config(monkeypatch, use_310p, model_type): + # Bypass config model loading, while retaining the real native predicates. + speculative = object.__new__(SpeculativeConfig) + speculative.method = "mtp" + speculative.enforce_eager = False + speculative.num_speculative_tokens = 1 + speculative.draft_model_config = SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)) + config = SimpleNamespace(speculative_config=speculative) + factory = init_speculator + if use_310p: + # Execute the actual route without installing global 310P patches. + path = Path(v2_spec_decode.__file__).resolve().parents[3] / "patch/worker/patch_spec_decode_310.py" + node = next(n for n in ast.parse(path.read_text()).body if isinstance(n, ast.FunctionDef)) + namespace: dict[str, Any] = {"ascend_spec_decode": v2_spec_decode, "_orig_init_speculator": init_speculator} + exec(compile(ast.Module(body=[node], type_ignores=[]), str(path), "exec"), namespace) + factory = namespace["_init_speculator_310p"] + if model_type == "step3p5_mtp": + module, name = "vllm_ascend.worker.v2.spec_decode.eagle.speculator", "AscendEagleSpeculator" + elif model_type == "gemma4_mtp": + module, name = "vllm_ascend.worker.v2.spec_decode.gemma4.speculator", "AscendGemma4Speculator" + elif use_310p: + module, name = "vllm_ascend._310p.worker.v2.spec_decode.mtp_speculator", "AscendMTPSpeculator310" + else: + module, name = "vllm_ascend.worker.v2.spec_decode.mtp.speculator", "AscendMTPSpeculator" + constructor = _mock_speculator_module(monkeypatch, module, name) + device = torch.device("cpu") + assert factory(config, device) is constructor.return_value + constructor.assert_called_once_with(config, device) diff --git a/tests/ut/worker/v2/test_speculator.py b/tests/ut/worker/v2/test_speculator.py index cb203af62e46..e88f6ce1f150 100644 --- a/tests/ut/worker/v2/test_speculator.py +++ b/tests/ut/worker/v2/test_speculator.py @@ -20,6 +20,7 @@ from vllm_ascend.attention.mla_v1 import AscendMLABackend from vllm_ascend.attention.sfa_v1 import AscendSFABackend from vllm_ascend.distributed.eplb.tensor_ops import make_allgather_mask +from vllm_ascend.utils import get_config_hash, vllm_version_is from vllm_ascend.worker.v2.input_batch import AscendInputBuffers from vllm_ascend.worker.v2.spec_decode.target_dependent_ar.speculator import ( AscendTargetDependentARSpeculator, @@ -169,7 +170,9 @@ def test_generate_draft_updates_attn_metadata(): speculator = AscendTargetDependentARSpeculator.__new__(AscendTargetDependentARSpeculator) attn_metadata = {"layer.0": MagicMock()} speculator._update_decode_attn_metadata = MagicMock() - with patch.object(TargetDependentARSpeculator, "_generate_draft") as parent_generate: + with patch.object( + TargetDependentARSpeculator, "_generate_draft" if vllm_version_is("0.33.0") else "_decode" + ) as parent_generate: # main always passes the per-batch step count; there is no fallback to # the configured count. speculator._generate_draft(2, 4, attn_metadata, None, None, num_speculative_steps=3) @@ -185,7 +188,7 @@ def test_generate_draft_skips_update_without_attn_metadata(): speculator = AscendTargetDependentARSpeculator.__new__(AscendTargetDependentARSpeculator) speculator.num_speculative_steps = 3 speculator._update_decode_attn_metadata = MagicMock() - with patch.object(TargetDependentARSpeculator, "_generate_draft"): + with patch.object(TargetDependentARSpeculator, "_generate_draft" if vllm_version_is("0.33.0") else "_decode"): speculator._generate_draft(2, 4, None, None, None) speculator._update_decode_attn_metadata.assert_not_called() @@ -749,7 +752,9 @@ def test_eplb_draft_mask_tracks_prefill_decode_and_dummy_replay(): speculator.input_buffers = SimpleNamespace(is_padding=torch.zeros(8, dtype=torch.bool)) speculator._eplb_source_is_padding = None speculator._eplb_dummy_run = False - speculator.speculative_config = SimpleNamespace(draft_model_config=SimpleNamespace(compute_hash=lambda: "draft")) + speculator.speculative_config = SimpleNamespace( + draft_model_config=SimpleNamespace(compute_hash=lambda: "draft", compile_factors=lambda: {"identity": "draft"}) + ) layer_state = SimpleNamespace() state = SimpleNamespace( allgather_record_tensor=torch.tensor(-1), @@ -757,6 +762,9 @@ def test_eplb_draft_mask_tracks_prefill_decode_and_dummy_replay(): "draft": SimpleNamespace(model=SimpleNamespace(moe_layers=[SimpleNamespace(eplb_state=layer_state)])) }, ) + state.model_states[get_config_hash(speculator.speculative_config.draft_model_config, cached=True)] = ( + state.model_states["draft"] + ) speculator.set_eplb_state(state) mask = layer_state.allgather_source_is_padding speculator.set_eplb_state(state) @@ -800,9 +808,11 @@ def parent(*args): assert speculator._lmhead_pad_size == num_reqs module = "vllm_ascend.worker.v2.spec_decode." + # #60989 renamed the native single-step draft entry point on main. + parent_method = "_generate_draft" if vllm_version_is("0.33.0") else "_decode" with ( patch(module + "lmhead_tp_utils.lmhead_tp_enable", return_value=True), - patch.object(TargetDependentARSpeculator, "_generate_draft", side_effect=parent) as generate, + patch.object(TargetDependentARSpeculator, parent_method, side_effect=parent) as generate, ): speculator._generate_draft(num_reqs, num_tokens, None, None, None) expected = [num_reqs, num_tokens, None, None, None, CUDAGraphMode.NONE] diff --git a/tests/ut/worker/v2/test_standalone_ar_speculator.py b/tests/ut/worker/v2/test_standalone_ar_speculator.py new file mode 100644 index 000000000000..49d686a6325d --- /dev/null +++ b/tests/ut/worker/v2/test_standalone_ar_speculator.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +import torch + +from vllm_ascend.utils import vllm_version_is +from vllm_ascend.worker.v2.spec_decode.standalone_ar import speculator as module + + +@pytest.mark.parametrize("step", [0, 1, 2]) +@pytest.mark.parametrize("target_bounds", [[7, 31], [63, 64]]) +def test_correction_token_is_counted_once_in_standalone_attention(step, target_bounds): + speculator = object.__new__(module.AscendStandaloneARSpeculator) + speculator.expanded_positions = torch.arange(8) + speculator.input_buffers = SimpleNamespace(positions=torch.arange(2)) + speculator.draft_is_prefilling = torch.tensor([True, False]) + speculator.vllm_config = SimpleNamespace(parallel_config=object(), model_config=SimpleNamespace(max_model_len=64)) + target = torch.tensor(target_bounds, dtype=torch.int32) + # #60989's main propose includes the correction before builder dispatch. + incoming = target.clone() if vllm_version_is("0.33.0") else target + 1 + original = incoming.clone() + qsl = np.array([0, 3, 8], dtype=np.int32) + desc, slots, metadata = object(), torch.arange(8), {"draft": object()} + + with ( + patch.object(module, "build_attn_metadata_wrapper", return_value=nullcontext()), + patch.object(module, "build_attn_metadata_factory", return_value=nullcontext()) as factory, + patch.object(module.StandaloneARSpeculator, "_build_attn_metadata", return_value=metadata) as parent, + ): + actual = speculator._build_attn_metadata(2, desc, qsl, incoming, step, slot_mappings=slots) + + expected = (target + 1).clamp(max=64) if step == 0 and vllm_version_is("0.33.0") else incoming + assert actual is metadata + torch.testing.assert_close(parent.call_args.args[3], expected) + torch.testing.assert_close(incoming, original) + torch.testing.assert_close(target, torch.tensor(target_bounds, dtype=torch.int32)) + assert parent.call_args.args[0] == 2 + assert parent.call_args.args[1] is desc + assert parent.call_args.args[2] is qsl + assert parent.call_args.args[4] == step + assert parent.call_args.args[7] is slots + expected_positions = speculator.expanded_positions if step == 0 else speculator.input_buffers.positions + assert factory.call_args.args[0] is expected_positions + assert factory.call_args.args[1] == 8 diff --git a/vllm_ascend/_310p/fused_moe/fused_moe.py b/vllm_ascend/_310p/fused_moe/fused_moe.py index 1c681ddc6637..5362ce06c6d5 100644 --- a/vllm_ascend/_310p/fused_moe/fused_moe.py +++ b/vllm_ascend/_310p/fused_moe/fused_moe.py @@ -167,6 +167,8 @@ def __init__( routed_output_transform=None, routed_scaling_factor=1, n_shared_experts: int = 0, + divide_shared_output_in_fp16=False, + routed_scale_mode=None, ): super().__init__( layer_name=layer_name, @@ -180,6 +182,8 @@ def __init__( routed_input_transform=routed_input_transform, routed_output_transform=routed_output_transform, routed_scaling_factor=routed_scaling_factor, + divide_shared_output_in_fp16=divide_shared_output_in_fp16, + routed_scale_mode=routed_scale_mode, ) gate_weight = getattr(gate, "weight", None) if isinstance(gate_weight, torch.Tensor) and gate_weight.dtype == torch.float16: diff --git a/vllm_ascend/_310p/worker/v2/model_runner.py b/vllm_ascend/_310p/worker/v2/model_runner.py index 67add61e7ba7..8f9cc451b7b6 100644 --- a/vllm_ascend/_310p/worker/v2/model_runner.py +++ b/vllm_ascend/_310p/worker/v2/model_runner.py @@ -54,7 +54,12 @@ from vllm_ascend._310p.worker.v2.states import Ascend310PRequestState from vllm_ascend.core.kv_cache_interface import get_storage_block_size from vllm_ascend.ops.rotary_embedding import update_cos_sin -from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, get_kv_cache_tensor_layers, shared_experts_calculation_stream +from vllm_ascend.utils import ( + ACL_FORMAT_FRACTAL_NZ, + get_kv_cache_tensor_layers, + shared_experts_calculation_stream, + vllm_version_is, +) from vllm_ascend.worker.v2.attn_utils import build_attn_state from vllm_ascend.worker.v2.model_runner import NPUModelRunner @@ -825,7 +830,7 @@ def initialize_kv_cache( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, self.decode_query_len, - use_v2_model_runner=True, + **({"use_v2_model_runner": True} if vllm_version_is("0.33.0") else {}), kv_cache_config=kv_cache_config, max_num_reqs=self.max_num_reqs, **{}, @@ -834,7 +839,9 @@ def initialize_kv_cache( self.vllm_config, self.device, cudagraph_mode, - self.decode_query_len, + self.decode_query_len + if vllm_version_is("0.33.0") + else self.vllm_config.get_decode_cudagraph_query_lens(self.decode_query_len), self, lora_capture_cases=self.lora_capture_cases, ) diff --git a/vllm_ascend/_310p/worker/v2/rope.py b/vllm_ascend/_310p/worker/v2/rope.py index 32507d275845..a050a16ac193 100644 --- a/vllm_ascend/_310p/worker/v2/rope.py +++ b/vllm_ascend/_310p/worker/v2/rope.py @@ -13,6 +13,7 @@ from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm_ascend._310p.worker.v2.states import Ascend310PStagedWriteTensor +from vllm_ascend.utils import vllm_version_is class Ascend310PRopeState: @@ -96,13 +97,20 @@ def prepare_positions_cpu( # Write directly into persistent host storage. # Writes directly into persistent output; no per-request # torch.arange allocation or intermediate CPU copy. - MRotaryEmbedding.get_next_input_positions_tensor( - out=self.positions_np, - out_offset=query_start, - mrope_position_delta=delta, - context_len=num_computed, - num_new_tokens=query_len, - ) + if vllm_version_is("0.33.0"): + MRotaryEmbedding.get_next_input_positions_tensor( + out=self.positions_np, + out_offset=query_start, + mrope_position_delta=delta, + context_len=num_computed, + num_new_tokens=query_len, + ) + else: + # Native #63373 removed the NumPy output-buffer helper. + # Keep its row broadcast and persistent host destination. + self.positions_np[:, query_start:query_end] = np.arange( + delta + num_computed, delta + num_computed + query_len, dtype=self.positions_np.dtype + ) # Keep the extra dummy column in the transfer. Slicing all rows but # only active columns makes a non-contiguous 2D view because storage diff --git a/vllm_ascend/_310p/worker/v2/spec_decode/mtp_speculator.py b/vllm_ascend/_310p/worker/v2/spec_decode/mtp_speculator.py index 9215141c1bc9..54d877b62847 100644 --- a/vllm_ascend/_310p/worker/v2/spec_decode/mtp_speculator.py +++ b/vllm_ascend/_310p/worker/v2/spec_decode/mtp_speculator.py @@ -34,6 +34,7 @@ from vllm_ascend._310p.ops.rotary_embedding import AscendRotaryEmbedding310 from vllm_ascend._310p.worker.v2.spec_utils import set_draft_step_host from vllm_ascend.attention.attention_v1 import AscendAttentionState +from vllm_ascend.utils import vllm_version_is from vllm_ascend.worker.v2.attn_utils import build_attn_metadata_factory from vllm_ascend.worker.v2.spec_decode.target_dependent_ar.speculator import ( AscendTargetDependentARSpeculator, @@ -274,17 +275,30 @@ def _generate_draft( """Skip Ascend metadata H2D under ACLGraph capture (pageable memcpy ban).""" # Call GPU AR generate_draft (sample + update_draft_inputs) without the # Ascend post-step ``seq_lens_cpu.copy_`` which is illegal while capturing. - TargetDependentARSpeculator._generate_draft( - self, - num_reqs, - num_tokens_padded, - attn_metadata, - slot_mappings, - num_tokens_across_dp, - cudagraph_runtime_mode, - # vLLM #60335 added `num_speculative_steps` on main. - num_speculative_steps, - ) + if vllm_version_is("0.33.0"): + TargetDependentARSpeculator._generate_draft( + self, + num_reqs, + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + # vLLM #60335 added `num_speculative_steps` on main. + num_speculative_steps, + ) + else: + TargetDependentARSpeculator._decode( + self, + num_reqs, + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + # vLLM #60335 added `num_speculative_steps` on main. + num_speculative_steps, + ) if attn_metadata is None or torch.npu.is_current_stream_capturing(): return self._update_decode_attn_metadata(attn_metadata, 1, num_reqs) @@ -293,6 +307,9 @@ def _set_draft_step(self, step: int) -> None: self.current_draft_step.fill_(step) set_draft_step_host(step) + # #60989 invokes _decode through the new shared AR loop. + _decode = _generate_draft + def _multi_step_decode( self, num_reqs: int, diff --git a/vllm_ascend/_310p/worker_310p.py b/vllm_ascend/_310p/worker_310p.py index c9ea1c157cd1..7f96cd84677b 100644 --- a/vllm_ascend/_310p/worker_310p.py +++ b/vllm_ascend/_310p/worker_310p.py @@ -156,7 +156,7 @@ def _init_device(self): # take current memory snapshot self.init_snapshot = MemorySnapshot(device=device) - self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization + self.requested_memory = self.init_snapshot.total_memory * self._get_gpu_memory_utilization() if is_rc_device(): self.init_snapshot.free_memory = psutil.virtual_memory().available logger.info_once("Root Complex (RC) mode: host and device memory are shared.") @@ -167,7 +167,7 @@ def _init_device(self): f"({GiB(self.init_snapshot.free_memory)}/" f"{GiB(self.init_snapshot.total_memory)} GiB) on startup " f"is less than desired GPU memory utilization " - f"({self.cache_config.gpu_memory_utilization}, " + f"({self._get_gpu_memory_utilization()}, " f"{GiB(self.requested_memory)} GiB). Decrease GPU memory " f"utilization or reduce GPU memory used by other processes." ) diff --git a/vllm_ascend/ascend_forward_context.py b/vllm_ascend/ascend_forward_context.py index b9a3b00fbe2a..dd8b31f282c2 100644 --- a/vllm_ascend/ascend_forward_context.py +++ b/vllm_ascend/ascend_forward_context.py @@ -21,6 +21,7 @@ from vllm_ascend.utils import ( has_layer_idx, is_moe_model, + vllm_version_is, ) # Dynamo constant-folds this like VLLM_USE_V2_MODEL_RUNNER. Sync from eager @@ -270,6 +271,32 @@ def set_mc2_tokens_capacity(vllm_config, max_num_reqs, uniform_decode_query_len) max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens elif vllm_config.compilation_config.cudagraph_capture_sizes: max_num_tokens = vllm_config.compilation_config.max_cudagraph_capture_size + compilation_config = vllm_config.compilation_config + decode_sizes = None + if not vllm_version_is("0.33.0"): + decode_sizes = compilation_config.decode_cudagraph_capture_request_sizes + if decode_sizes is not None: + # Main's shared token union also includes independent one-token + # draft graphs up to max_num_seqs (#63406). Keep the communication + # threshold tied to target graphs: wider draft batches use the + # existing ALLTOALL fallback instead of enlarging the MC2 profile. + # This helper is only available with the split-grid main API. + from vllm.config.compilation import decode_cudagraph_capture_shapes + + decode_max_tokens = min( + vllm_config.scheduler_config.max_num_batched_tokens, + compilation_config.max_decode_cudagraph_capture_tokens + or vllm_config.scheduler_config.max_num_batched_tokens, + ) + target_shapes = decode_cudagraph_capture_shapes( + decode_sizes, + vllm_config.get_decode_cudagraph_query_lens(uniform_decode_query_len), + decode_max_tokens, + ) + max_num_tokens = max( + compilation_config.max_prefill_cudagraph_capture_tokens or 0, + max((num_reqs * query_len for num_reqs, query_len in target_shapes), default=0), + ) else: max_num_tokens = max_num_reqs * uniform_decode_query_len tp_size = vllm_config.parallel_config.tensor_parallel_size diff --git a/vllm_ascend/attention/mla_v1.py b/vllm_ascend/attention/mla_v1.py index dd5b21075ef4..69388f4d6dd5 100644 --- a/vllm_ascend/attention/mla_v1.py +++ b/vllm_ascend/attention/mla_v1.py @@ -88,12 +88,21 @@ enable_sfa, is_pd_decode_recompute_scheduler_enabled, maybe_trans_nz, + vllm_version_is, weak_ref_tensors, ) BUILD_METADATA_STEP_PREFILL = 0 BUILD_METADATA_STEP_DECODE = 1 + +def _uses_unpadded_drafter_batch(speculative_config) -> bool: + if vllm_version_is("0.33.0"): + return speculative_config.disable_padded_drafter_batch + # MRv1's optional unpadded drafter was removed in native #63373. + return False + + # Exclusive batch * K limit of the fused op with perm_x1=(1, 0, 2). TRANSPOSE_BMM_MAX_SUPPORTED_DIM = 65536 @@ -715,9 +724,8 @@ def get_block_table_size(self, common_attn_metadata: AscendCommonAttentionMetada if build_metadata_step == BUILD_METADATA_STEP_PREFILL: # If graph_pad_size > -1, mean is running in fullgraph mode. # NOTE: Maybe this block_table change can be removed when graph_pad_size > 1. - if ( - self.graph_pad_size > common_attn_metadata.num_reqs - and self.speculative_config.disable_padded_drafter_batch + if self.graph_pad_size > common_attn_metadata.num_reqs and _uses_unpadded_drafter_batch( + self.speculative_config ): return self.graph_pad_size return common_attn_metadata.num_reqs @@ -786,12 +794,12 @@ def build_decode_metadata( # NOTE: MTP full graph is incompatible with context parallelism. # NOTE: Maybe this block_table change can be removed when graph_pad_size > 1. - if self.graph_pad_size > self.num_decodes and self.speculative_config.disable_padded_drafter_batch: + if self.graph_pad_size > self.num_decodes and _uses_unpadded_drafter_batch(self.speculative_config): self.block_table = self.block_table[: self.graph_pad_size, ...] seq_lens_list = self.seq_lens.tolist() if self.graph_pad_size > num_reqs: - if self.speculative_config.disable_padded_drafter_batch: + if _uses_unpadded_drafter_batch(self.speculative_config): num_reqs_pad_size = self.graph_pad_size - num_reqs actual_seq_lengths_q = self.pad_actual_seq_len_q_mtp_disable_pad( num_reqs_pad_size, num_reqs, actual_seq_lengths_q @@ -1140,7 +1148,7 @@ def update_graph_params( actual_seq_lengths = layer_metadata.decode.actual_seq_lengths_q block_table = layer_metadata.decode.block_table # TODO: This is a hack and should be fixed in the future. - if speculative_config.disable_padded_drafter_batch: + if _uses_unpadded_drafter_batch(speculative_config): block_table = block_table[: len(actual_seq_lengths)] seq_lens_list = seq_lens_list + [0] * (len(actual_seq_lengths) - len(seq_lens_list)) else: diff --git a/vllm_ascend/attention/sfa_v1.py b/vllm_ascend/attention/sfa_v1.py index 8c9245de89f4..0ca894ca02e0 100644 --- a/vllm_ascend/attention/sfa_v1.py +++ b/vllm_ascend/attention/sfa_v1.py @@ -17,6 +17,7 @@ from vllm.v1.attention.backend import ( AttentionBackend, # type: ignore AttentionCGSupport, + AttentionMetadataBuilder, MLAAttentionImpl, ) from vllm.v1.kv_cache_interface import AttentionSpec @@ -77,6 +78,7 @@ is_mtp_layer, is_rl_weight_update_enabled, maybe_trans_nz, + vllm_version_is, ) from vllm_ascend.weight_switch import ParallelOProjection, WeightSwitchConfig @@ -524,14 +526,22 @@ def __init__( metadata_cls: type[AscendSFAMetadata] | None = None, supports_dcp_with_varlen: bool = False, ): - super().__init__( - kv_cache_spec, - layer_names, - vllm_config, - device, - metadata_cls if metadata_cls is not None else AscendSFAMetadata, - supports_dcp_with_varlen, - ) + if vllm_version_is("0.33.0"): + super().__init__( + kv_cache_spec, + layer_names, + vllm_config, + device, + metadata_cls if metadata_cls is not None else AscendSFAMetadata, + supports_dcp_with_varlen, + ) + else: + # SFA builds its own sparse metadata and never uses the dense MLA + # prefill workspace/backend. Native sparse attention has no dense + # backend to clone in MLACommonMetadataBuilder.__init__. + AttentionMetadataBuilder.__init__(self, kv_cache_spec, layer_names, vllm_config, device) + self.metadata_cls = metadata_cls if metadata_cls is not None else AscendSFAMetadata + self.model_config = vllm_config.model_config # Match the logical block size selected for BlockTable. self.kernel_block_size = select_common_block_size(kv_cache_spec.block_size, [AscendSFABackend]) @@ -752,6 +762,11 @@ class AscendSFAImpl(MLAAttentionImpl): understand this class """ + # Main validates the sparse contract before our custom forward() runs. + # SFA consumes top-k indices itself and has no native dense-MHA entrypoint. + is_sparse: ClassVar[bool] = not vllm_version_is("0.33.0") + supports_dense_mha_prefill: bool = vllm_version_is("0.33.0") + # ``W_UV``/``W_UK_T`` are injected during ``process_weights_after_loading`` # through ``replace_parameter``, which is ``setattr``-based and therefore # invisible to static analysis. These declarations are what let mypy resolve diff --git a/vllm_ascend/compilation/acl_graph.py b/vllm_ascend/compilation/acl_graph.py index b307581dac73..1a8a98399846 100644 --- a/vllm_ascend/compilation/acl_graph.py +++ b/vllm_ascend/compilation/acl_graph.py @@ -14,7 +14,6 @@ import vllm.envs as envs from vllm.compilation.counter import compilation_counter from vllm.compilation.cuda_graph import CUDAGraphOptions -from vllm.compilation.monitor import validate_cudagraph_capturing_enabled from vllm.config import CUDAGraphMode, VllmConfig from vllm.forward_context import BatchDescriptor, get_forward_context, is_forward_context_available from vllm.logger import logger @@ -23,7 +22,7 @@ from vllm_ascend.ascend_config import get_ascend_config from vllm_ascend.ascend_forward_context import _EXTRA_CTX -from ..utils import super_kernel_scope, use_updatable_graph, weak_ref_tensors +from ..utils import super_kernel_scope, use_updatable_graph, vllm_version_is, weak_ref_tensors from .updatable_graph import ( ContextSource, SharedSource, @@ -31,6 +30,17 @@ ) _acl_graph_wrappers: weakref.WeakSet[Any] = weakref.WeakSet() + +if vllm_version_is("0.33.0"): + from vllm.compilation import monitor + + +def validate_cudagraph_capturing_enabled() -> None: + """Retain release's capture guard until v0.33.0 support ends.""" + if vllm_version_is("0.33.0"): + monitor.validate_cudagraph_capturing_enabled() + + _STREAM_RESOURCE_ERROR_CODE = "207008" _STREAM_RESOURCE_ERROR_MARKERS = ( "insufficient_stream_resources", diff --git a/vllm_ascend/compilation/compiler_interface.py b/vllm_ascend/compilation/compiler_interface.py index 23db1127f442..e0c56501d050 100644 --- a/vllm_ascend/compilation/compiler_interface.py +++ b/vllm_ascend/compilation/compiler_interface.py @@ -23,6 +23,7 @@ import torch import torch.fx as fx +import torch_npu from npugraph_ex._acl_concrete_graph.static_kernel import _set_static_kernel_blacklist from torch._dynamo.backends.common import aot_autograd from torch._inductor.compile_fx import graph_returns_tuple, make_graph_return_tuple @@ -35,7 +36,7 @@ from vllm.utils.math_utils import round_up from vllm_ascend.ascend_config import AscendCompilationConfig, get_ascend_config -from vllm_ascend.utils import COMPILATION_PASS_KEY +from vllm_ascend.utils import COMPILATION_PASS_KEY, vllm_version_is def compile_fx(graph: GraphModule, example_inputs: list, inner_compile: Callable, decompositions: dict) -> Callable: @@ -272,13 +273,9 @@ class AscendCompiler(CompilerInterface): name = "AscendCompiler" - # TODO(wxs): add passes related to compilation in compute_hash - def compute_hash(self, vllm_config: VllmConfig) -> str: + def compile_factors(self, vllm_config: VllmConfig) -> dict[str, Any]: self.vllm_config = vllm_config ascend_compilation_config = get_ascend_config().ascend_compilation_config - from hashlib import sha256 - - import torch_npu factors = { "torch_npu_version": torch_npu.__version__, @@ -287,7 +284,15 @@ def compute_hash(self, vllm_config: VllmConfig) -> str: "enable_super_kernel": ascend_compilation_config.enable_super_kernel, } logger.info("AscendCompiler hash factors: %s", factors) - return sha256(str(factors).encode(), usedforsecurity=False).hexdigest()[:10] + return factors + + if vllm_version_is("0.33.0"): + + def compute_hash(self, vllm_config: VllmConfig) -> str: + # Main rejects subclasses defining compute_hash (#65652). + from hashlib import sha256 + + return sha256(str(self.compile_factors(vllm_config)).encode(), usedforsecurity=False).hexdigest()[:10] def initialize_cache(self, cache_dir: str, disable_cache: bool = False, prefix: str = "") -> None: """Initialize cache settings and runtime state before compiling or loading a graph.""" diff --git a/vllm_ascend/core/dyntra_lb_scheduler.py b/vllm_ascend/core/dyntra_lb_scheduler.py index 04bbf7a44177..9794bbab5fbb 100644 --- a/vllm_ascend/core/dyntra_lb_scheduler.py +++ b/vllm_ascend/core/dyntra_lb_scheduler.py @@ -30,6 +30,8 @@ from vllm.v1.utils import record_function_or_nullcontext from vllm_ascend.ascend_config import DyntraLBConfig +from vllm_ascend.core.scheduler_output import make_cached_request_data +from vllm_ascend.utils import vllm_version_is if TYPE_CHECKING: from vllm.v1.core.sched.scheduler import Scheduler as _SchedulerBase @@ -1064,16 +1066,8 @@ def skip_request(from_queue) -> None: # len(self.running). assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len(scheduled_running_reqs) <= len(self.running) - # Get the longest common prefix among all requests in the running queue. - # This can be potentially used for cascade attention. - num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) - with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): - if self.running: - any_request_id = self.running[0].request_id - num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) - # Construct the scheduler output. - if self.use_v2_model_runner: + if not vllm_version_is("0.33.0") or self.use_v2_model_runner: scheduled_new_reqs.extend(scheduled_resumed_reqs) scheduled_resumed_reqs.clear() new_reqs_data = [ @@ -1096,16 +1090,20 @@ def skip_request(from_queue) -> None: ] with record_function_or_nullcontext("schedule: make_cached_request_data"): - cached_reqs_data = self._make_cached_request_data( + cached_reqs_data, schedule_metadata = make_cached_request_data( + self, scheduled_running_reqs, scheduled_resumed_reqs, num_scheduled_tokens, scheduled_spec_decode_tokens, req_to_new_blocks, + new_reqs_data, + scheduled_encoder_inputs, + encoder_compute_budget, ) # Record the request ids that were scheduled in this step (MRV1-only). - if not self.use_v2_model_runner: + if vllm_version_is("0.33.0") and not self.use_v2_model_runner: self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) @@ -1142,9 +1140,11 @@ def skip_request(from_queue) -> None: if self.dynamic_sd_lookup is not None and len(num_scheduled_tokens) > 0: num_spec_tokens_to_schedule = self.dynamic_sd_lookup[len(num_scheduled_tokens)] - scheduled_encoder_input_stats = None - if self.log_stats and self.observability_config.enable_logging_iteration_details: - scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats(scheduled_encoder_inputs) + if vllm_version_is("0.33.0"): + scheduled_encoder_input_stats = None + if self.log_stats and self.observability_config.enable_logging_iteration_details: + scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats(scheduled_encoder_inputs) + schedule_metadata["scheduled_encoder_input_stats"] = scheduled_encoder_input_stats scheduler_output_kwargs = dict( scheduled_new_reqs=new_reqs_data, @@ -1153,8 +1153,7 @@ def skip_request(from_queue) -> None: total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, - scheduled_encoder_input_stats=scheduled_encoder_input_stats, - num_common_prefix_blocks=num_common_prefix_blocks, + **schedule_metadata, preempted_req_ids=self.reset_preempted_req_ids, # finished_req_ids is an existing state in the scheduler, # instead of being newly scheduled in this step. diff --git a/vllm_ascend/core/recompute_scheduler.py b/vllm_ascend/core/recompute_scheduler.py index ad1a723b54eb..b89f8a939b04 100644 --- a/vllm_ascend/core/recompute_scheduler.py +++ b/vllm_ascend/core/recompute_scheduler.py @@ -46,7 +46,8 @@ DyntraLBPolicyMixin, print_scheduler_summary, ) -from vllm_ascend.utils import get_ascend_config +from vllm_ascend.core.scheduler_output import make_cached_request_data +from vllm_ascend.utils import get_ascend_config, vllm_version_is def _normalize_encoder_schedule_result( @@ -901,19 +902,11 @@ def skip_request(from_queue) -> None: # len(self.running). assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len(scheduled_running_reqs) <= len(self.running) - # Get the longest common prefix among all requests in the running queue. - # This can be potentially used for cascade attention. - num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) - with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): - if self.running: - any_request_id = self.running[0].request_id - num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) - # Construct the scheduler output. new_request_kwargs = { "uses_mrope": self.model_uses_mrope, } - if self.use_v2_model_runner: + if not vllm_version_is("0.33.0") or self.use_v2_model_runner: scheduled_new_reqs.extend(scheduled_resumed_reqs) scheduled_resumed_reqs.clear() new_reqs_data = [ @@ -936,16 +929,20 @@ def skip_request(from_queue) -> None: ] with record_function_or_nullcontext("schedule: make_cached_request_data"): - cached_reqs_data = self._make_cached_request_data( + cached_reqs_data, schedule_metadata = make_cached_request_data( + self, scheduled_running_reqs, scheduled_resumed_reqs, num_scheduled_tokens, scheduled_spec_decode_tokens, req_to_new_blocks, + new_reqs_data, + scheduled_encoder_inputs, + encoder_compute_budget, ) # Record the request ids that were scheduled in this step (MRV1-only). - if not self.use_v2_model_runner: + if vllm_version_is("0.33.0") and not self.use_v2_model_runner: self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) @@ -977,9 +974,11 @@ def skip_request(from_queue) -> None: if self.dynamic_sd_lookup is not None and len(num_scheduled_tokens) > 0: num_spec_tokens_to_schedule = self.dynamic_sd_lookup[len(num_scheduled_tokens)] - scheduled_encoder_input_stats = None - if self.log_stats and self.observability_config.enable_logging_iteration_details: - scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats(scheduled_encoder_inputs) + if vllm_version_is("0.33.0"): + scheduled_encoder_input_stats = None + if self.log_stats and self.observability_config.enable_logging_iteration_details: + scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats(scheduled_encoder_inputs) + schedule_metadata["scheduled_encoder_input_stats"] = scheduled_encoder_input_stats scheduler_output = RecomputeSchedulerOutput( scheduled_new_reqs=new_reqs_data, @@ -988,8 +987,7 @@ def skip_request(from_queue) -> None: total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, - scheduled_encoder_input_stats=scheduled_encoder_input_stats, - num_common_prefix_blocks=num_common_prefix_blocks, + **schedule_metadata, preempted_req_ids=self.reset_preempted_req_ids, # finished_req_ids is an existing state in the scheduler, # instead of being newly scheduled in this step. @@ -1000,7 +998,11 @@ def skip_request(from_queue) -> None: new_block_ids_to_zero=self._get_new_block_ids_to_zero(), kv_cache_block_copies=pending_kv_cache_block_copies, num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, - ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(), + **( + {"ec_manager_metadata": self.encoder_cache_manager.get_manager_metadata()} + if vllm_version_is("0.33.0") + else {} + ), recomputed_reqs=self._recomputed_reqs or None, has_sync_kv_loads=has_sync_kv_loads, kv_connector_block_state=kv_connector_block_state, diff --git a/vllm_ascend/core/scheduler_output.py b/vllm_ascend/core/scheduler_output.py new file mode 100644 index 000000000000..b4e774a0b0ef --- /dev/null +++ b/vllm_ascend/core/scheduler_output.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +"""vLLM release/main output construction shared by copied schedulers.""" + +from typing import Any + +from vllm_ascend.utils import vllm_version_is + +if not vllm_version_is("0.33.0"): + from vllm.v1.core.sched.output import IterationDetails + + +def make_cached_request_data( + scheduler: Any, + running_reqs: list, + resumed_reqs: list, + num_scheduled_tokens: dict[str, int], + spec_decode_tokens: dict[str, list[int]], + req_to_new_blocks: dict, + new_reqs_data: list, + scheduled_encoder_inputs: dict[str, list[int]], + encoder_compute_budget: int, +) -> tuple[Any, dict[str, Any]]: + """Keep release payloads and populate main's scheduling-time statistics. + + Resumed requests are already included in new_reqs_data on main. Encoder + counts use the final scheduled set, so preempted inputs are not counted. + """ + if vllm_version_is("0.33.0"): + return scheduler._make_cached_request_data( + running_reqs, resumed_reqs, num_scheduled_tokens, spec_decode_tokens, req_to_new_blocks + ), {} + + details = None + if ( + scheduler.log_stats and scheduler.observability_config.enable_logging_iteration_details + ) or scheduler.vllm_config.profiler_config.profiler is not None: + details = IterationDetails() + cached = scheduler._make_cached_request_data( + running_reqs, num_scheduled_tokens, spec_decode_tokens, req_to_new_blocks, iteration_details=details + ) + if details is not None: + details.num_prefill_requests = len(num_scheduled_tokens) - details.num_decode_requests + details.num_prefill_tokens = sum(num_scheduled_tokens.values()) - details.num_decode_tokens + details.num_encoder_inputs = sum(len(inputs) for inputs in scheduled_encoder_inputs.values()) + details.num_encoder_output_tokens = scheduler.max_num_encoder_input_tokens - encoder_compute_budget + return cached, { + "iteration_details": details, + "encoder_input_updates": scheduler._get_encoder_input_updates( + num_scheduled_tokens, new_reqs_data, scheduled_encoder_inputs + ), + } diff --git a/vllm_ascend/core/scheduler_profiling_chunk.py b/vllm_ascend/core/scheduler_profiling_chunk.py index 16806fa86a0f..ccf7a4cf586a 100644 --- a/vllm_ascend/core/scheduler_profiling_chunk.py +++ b/vllm_ascend/core/scheduler_profiling_chunk.py @@ -46,6 +46,8 @@ from vllm.v1.utils import record_function_or_nullcontext from vllm_ascend.core.profiling_chunk_predictor import ProfilingChunkManager +from vllm_ascend.core.scheduler_output import make_cached_request_data +from vllm_ascend.utils import vllm_version_is class ProfilingChunkScheduler(Scheduler): @@ -1036,14 +1038,6 @@ def skip_request(from_queue) -> None: # len(self.running). assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len(scheduled_running_reqs) <= len(self.running) - # Get the longest common prefix among all requests in the running queue. - # This can be potentially used for cascade attention. - num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) - with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): - if self.running: - any_request_id = self.running[0].request_id - num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) - # Construct the scheduler output. scheduled_new_reqs.extend(scheduled_resumed_reqs) scheduled_resumed_reqs.clear() @@ -1058,12 +1052,16 @@ def skip_request(from_queue) -> None: ] with record_function_or_nullcontext("schedule: make_cached_request_data"): - cached_reqs_data = self._make_cached_request_data( + cached_reqs_data, schedule_metadata = make_cached_request_data( + self, scheduled_running_reqs, scheduled_resumed_reqs, num_scheduled_tokens, scheduled_spec_decode_tokens, req_to_new_blocks, + new_reqs_data, + scheduled_encoder_inputs, + encoder_compute_budget, ) # Mamba "align" boundary states must be handed off with exact block ids; @@ -1098,9 +1096,11 @@ def skip_request(from_queue) -> None: if self.dynamic_sd_lookup is not None and len(num_scheduled_tokens) > 0: num_spec_tokens_to_schedule = self.dynamic_sd_lookup[len(num_scheduled_tokens)] - scheduled_encoder_input_stats = None - if self.log_stats and self.observability_config.enable_logging_iteration_details: - scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats(scheduled_encoder_inputs) + if vllm_version_is("0.33.0"): + scheduled_encoder_input_stats = None + if self.log_stats and self.observability_config.enable_logging_iteration_details: + scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats(scheduled_encoder_inputs) + schedule_metadata["scheduled_encoder_input_stats"] = scheduled_encoder_input_stats scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, @@ -1109,8 +1109,7 @@ def skip_request(from_queue) -> None: total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, - scheduled_encoder_input_stats=scheduled_encoder_input_stats, - num_common_prefix_blocks=num_common_prefix_blocks, + **schedule_metadata, preempted_req_ids=self.reset_preempted_req_ids, # finished_req_ids is an existing state in the scheduler, # instead of being newly scheduled in this step. @@ -1123,7 +1122,11 @@ def skip_request(from_queue) -> None: kv_cache_block_copies=pending_kv_cache_block_copies, kv_connector_block_state=kv_connector_block_state, num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, - ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(), + **( + {"ec_manager_metadata": self.encoder_cache_manager.get_manager_metadata()} + if vllm_version_is("0.33.0") + else {} + ), ) # NOTE(Kuntai): this function is designed for multiple purposes: diff --git a/vllm_ascend/distributed/ec_transfer/ec_connector/mooncake/store_backend.py b/vllm_ascend/distributed/ec_transfer/ec_connector/mooncake/store_backend.py index ef46e63b937f..d18a525e5284 100644 --- a/vllm_ascend/distributed/ec_transfer/ec_connector/mooncake/store_backend.py +++ b/vllm_ascend/distributed/ec_transfer/ec_connector/mooncake/store_backend.py @@ -31,6 +31,7 @@ from vllm_ascend.distributed.ec_transfer.ec_connector.mooncake.memory import AscendProducerMemoryPool from vllm_ascend.distributed.ec_transfer.ec_connector.mooncake.store_client import AscendMooncakeEmbeddingStoreClient from vllm_ascend.distributed.ec_transfer.ec_connector.mooncake.transfer import AscendMooncakeTransfer +from vllm_ascend.utils import vllm_version_is def use_store_independent_te() -> bool: @@ -115,7 +116,10 @@ def create_ascend_mooncake_embedding_store_client( except ImportError as error: raise ImportError("Install mooncake with Store support to enable cross_encoder_cache.") from error - from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils + if vllm_version_is("0.33.0"): + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils + else: + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake_store import rdma_utils store = MooncakeDistributedStore() local_hostname = rdma_utils.get_requester_local_hostname(get_ip()) diff --git a/vllm_ascend/distributed/eplb/eplb_state.py b/vllm_ascend/distributed/eplb/eplb_state.py index e3cb87bd5042..7f8ddfd41e29 100644 --- a/vllm_ascend/distributed/eplb/eplb_state.py +++ b/vllm_ascend/distributed/eplb/eplb_state.py @@ -36,6 +36,7 @@ count_allgather_tokens, make_allgather_mask, ) +from vllm_ascend.utils import get_config_hash ASYNC_EPLB_CYCLE_COMMITTED_LOG = "Ascend async EPLB cycle committed" _EPLB_CLOSE_TIMEOUT_SECONDS = 300 @@ -314,10 +315,10 @@ def add_model(self, model, model_config, model_name: str | None = None) -> None: # refresh_model_routing_tables needs the previous placement to notice. # Seeding here rather than on first use also covers the first commit, # which already carries the new mapping by the time the tables refresh. - model_state = self.model_states[model_config.compute_hash()] + model_state = self.model_states[get_config_hash(model_config, cached=True)] model_state._allgather_physical_map = model_state.physical_to_logical_map.clone() if self.uses_custom_load_stats: - self._initialize_load_stats_state(self.model_states[model_config.compute_hash()]) + self._initialize_load_stats_state(model_state) def _propagate_shared_tensors(self, model, num_unpadded_tokens_tensors) -> None: super()._propagate_shared_tensors(model, num_unpadded_tokens_tensors) @@ -726,7 +727,7 @@ def update_mapping(self, model_config, expanded_physical_to_logical: torch.Tenso super().update_mapping(model_config, expanded_physical_to_logical) # Verified main updates a live state during elastic EP. Refresh the # Ascend replica table in place so captured graphs use the new mapping. - refresh_model_routing_tables(self.model_states[model_config.compute_hash()]) + refresh_model_routing_tables(self.model_states[get_config_hash(model_config, cached=True)]) @classmethod def from_mapping( diff --git a/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/mooncake_layerwise.py b/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/mooncake_layerwise.py index b7a7292d6248..a7493584d574 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/mooncake_layerwise.py +++ b/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/backend/mooncake_layerwise.py @@ -19,6 +19,7 @@ block_hash_to_str, get_block_hashes, ) +from vllm_ascend.utils import get_config_hash if TYPE_CHECKING: from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.layerwise.pool_worker import LayerwiseKVPoolWorker @@ -97,11 +98,11 @@ def layerwise_topology_namespace(vllm_config: Any, kv_cache_config: Any = None) identity: tuple[Any, ...] = ( parallel.tensor_parallel_size, partitions, - model.compute_hash(), - cache.compute_hash(), + get_config_hash(model), + get_config_hash(cache), cache.cache_dtype, cache.block_size, - speculative.compute_hash() if speculative is not None else None, + get_config_hash(speculative) if speculative is not None else None, ) if kvpp_namespace: identity += (kvpp_namespace,) diff --git a/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/lookup_key.py b/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/lookup_key.py index 52cf7df4762e..5fc50b678bfa 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/lookup_key.py +++ b/vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/lookup_key.py @@ -1,22 +1,35 @@ from __future__ import annotations import threading +from importlib import import_module from typing import TYPE_CHECKING import vllm.envs as envs import zmq from vllm.config import VllmConfig -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import ( - LOOKUP_MSG, - RESET_MSG, - RESP_ERR, - RESP_OK, -) from vllm.logger import logger from vllm.utils.network_utils import make_zmq_socket from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder +from vllm_ascend.utils import vllm_version_is + +if TYPE_CHECKING or not vllm_version_is("0.33.0"): + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake_store.protocol import ( + LOOKUP_MSG, + RESET_MSG, + RESP_ERR, + RESP_OK, + ) +else: + # Static CI checks against the pinned main tree; the release module is + # loaded at runtime only, preserving its own wire constants. + _legacy_protocol = import_module("vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol") + LOOKUP_MSG = _legacy_protocol.LOOKUP_MSG + RESET_MSG = _legacy_protocol.RESET_MSG + RESP_ERR = _legacy_protocol.RESP_ERR + RESP_OK = _legacy_protocol.RESP_OK + if TYPE_CHECKING: from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.pool_worker import BaseKVPoolWorker diff --git a/vllm_ascend/distributed/kv_transfer/kv_pool/kv_offload/native/cpu_npu.py b/vllm_ascend/distributed/kv_transfer/kv_pool/kv_offload/native/cpu_npu.py index e639b58cd368..252c28d5f210 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_pool/kv_offload/native/cpu_npu.py +++ b/vllm_ascend/distributed/kv_transfer/kv_pool/kv_offload/native/cpu_npu.py @@ -8,6 +8,7 @@ from collections import deque from collections.abc import Callable from dataclasses import dataclass +from typing import TYPE_CHECKING import torch from vllm.logger import logger @@ -25,6 +26,11 @@ from vllm.v1.kv_offload.cpu.gpu_worker import compute_sub_block_ptrs from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion +from vllm_ascend.utils import vllm_version_is + +if TYPE_CHECKING: + from vllm.v1.kv_offload.base import DevicePointers + # Direction codes shared with csrc/torch_binding.cpp::swap_blocks_batch. DIRECTION_H2D = 0 DIRECTION_D2H = 1 @@ -107,6 +113,58 @@ def transfer_async( src_spec: LoadStoreSpec, dst_spec: LoadStoreSpec, ) -> bool: + if vllm_version_is("0.33.0"): + batch_src, batch_dst, batch_sizes, num_copy_ops, num_transfer_bytes = self._legacy_descriptors( + src_spec, dst_spec + ) + else: + device_ptrs = src_spec if self.npu_to_cpu else dst_spec + cpu_spec = dst_spec if self.npu_to_cpu else src_spec + batch_src, batch_dst, batch_sizes, num_copy_ops, num_transfer_bytes = self._pointer_descriptors( + device_ptrs, cpu_spec + ) + src = batch_src[:num_copy_ops] + dst = batch_dst[:num_copy_ops] + sizes = batch_sizes[:num_copy_ops] + + stream = self._stream_pool.pop() if self._stream_pool else torch.npu.Stream() + start_event = self._event_pool.pop() if self._event_pool else torch.npu.Event(enable_timing=True) + end_event = self._event_pool.pop() if self._event_pool else torch.npu.Event(enable_timing=True) + + if self.npu_to_cpu: + # KV writes happen on the model stream. D2H must not observe a + # partially written block. + stream.wait_stream(torch.npu.current_stream()) + if self._transfers: + stream.wait_event(self._transfers[-1].end_event) + + with torch.npu.stream(stream): + start_event.record(stream) + if num_copy_ops > 0: + torch.ops._C_ascend.swap_blocks_batch( + src, + dst, + sizes, + self.direction, + ) + end_event.record(stream) + + self._transfer_events[job_id] = end_event + self._transfers.append( + Transfer( + job_id=job_id, + stream=stream, + start_event=start_event, + end_event=end_event, + num_bytes=num_transfer_bytes, + batch_src=batch_src, + batch_dst=batch_dst, + batch_sizes=batch_sizes, + ) + ) + return True + + def _legacy_descriptors(self, src_spec: LoadStoreSpec, dst_spec: LoadStoreSpec): assert isinstance(src_spec, BlockIDsLoadStoreSpec) assert isinstance(dst_spec, BlockIDsLoadStoreSpec) @@ -193,42 +251,59 @@ def transfer_async( assert dst_offset == len(dst_blocks) assert op_idx == num_copy_ops - stream = self._stream_pool.pop() if self._stream_pool else torch.npu.Stream() - start_event = self._event_pool.pop() if self._event_pool else torch.npu.Event(enable_timing=True) - end_event = self._event_pool.pop() if self._event_pool else torch.npu.Event(enable_timing=True) - - if self.npu_to_cpu: - # KV writes happen on the model stream. D2H must not observe a - # partially written block. - stream.wait_stream(torch.npu.current_stream()) - if self._transfers: - stream.wait_event(self._transfers[-1].end_event) + return batch_src, batch_dst, batch_sizes, num_copy_ops, num_transfer_bytes - with torch.npu.stream(stream): - start_event.record(stream) - if num_copy_ops > 0: - torch.ops._C_ascend.swap_blocks_batch( - src, - dst, - sizes, - self.direction, - ) - end_event.record(stream) + def _pointer_descriptors(self, device_ptrs: DevicePointers, cpu_spec: LoadStoreSpec): + """Translate #56076's resolved device addresses into NPU DMA descriptors. - self._transfer_events[job_id] = end_event - self._transfers.append( - Transfer( - job_id=job_id, - stream=stream, - start_event=start_event, - end_event=end_event, - num_bytes=num_transfer_bytes, - batch_src=batch_src, - batch_dst=batch_dst, - batch_sizes=batch_sizes, - ) + Keep group/ref/block order, including empty groups and partial chunks. + CPU chunks retain the original physical stride; copy sizes are unpadded. + """ + assert isinstance(cpu_spec, BlockIDsLoadStoreSpec) + cpu_blocks = cpu_spec.block_ids + num_copy_ops = len(device_ptrs.ptrs) + batch_src, batch_dst, batch_sizes = ( + self._buffer_pool.pop() if self._buffer_pool else _new_descriptor_buffers(num_copy_ops) ) - return True + if batch_src.numel() < num_copy_ops: + batch_src, batch_dst, batch_sizes = _new_descriptor_buffers(num_copy_ops) + all_src = batch_src[:num_copy_ops].numpy() + all_dst = batch_dst[:num_copy_ops].numpy() + all_sizes = batch_sizes[:num_copy_ops].numpy() + npu_out, cpu_out = (all_src, all_dst) if self.npu_to_cpu else (all_dst, all_src) + npu_out[:] = device_ptrs.ptrs + all_sizes[:] = device_ptrs.sizes + blocks_per_chunk = self.dst_blocks_per_chunk if self.npu_to_cpu else self.src_blocks_per_chunk + cpu_offset = 0 + op_offset = 0 + assert len(device_ptrs.group_block_counts) == len(self.kv_cache_groups_data_refs) + for group_size, n_refs, block_idx, refs in zip( + device_ptrs.group_block_counts, + device_ptrs.group_data_ref_counts, + device_ptrs.block_indices, + self.kv_cache_groups_data_refs, + ): + assert n_refs == len(refs) + if group_size == 0: + continue + skip = block_idx % blocks_per_chunk + num_chunks = cdiv(group_size + skip, blocks_per_chunk) + group_blocks = cpu_blocks[cpu_offset : cpu_offset + num_chunks] + assert len(group_blocks) == num_chunks + for ref in refs: + end = op_offset + group_size + compute_sub_block_ptrs( + group_blocks, + blocks_per_chunk, + cpu_out[op_offset:end], + (self.dst_tensors if self.npu_to_cpu else self.src_tensors)[ref.tensor_idx], + skip_count=skip, + ) + op_offset = end + cpu_offset += num_chunks + assert cpu_offset == len(cpu_blocks) + assert op_offset == num_copy_ops + return batch_src, batch_dst, batch_sizes, num_copy_ops, int(device_ptrs.sizes.sum()) def get_finished(self) -> list[TransferResult]: results: list[TransferResult] = [] @@ -359,21 +434,21 @@ def __init__( def submit_store( self, job_id: int, - src_spec: GPULoadStoreSpec, + device_ptrs: DevicePointers | GPULoadStoreSpec, dst_spec: LoadStoreSpec, ) -> bool: - return self._store_handler.transfer_async(job_id, src_spec, dst_spec) + return self._store_handler.transfer_async(job_id, device_ptrs, dst_spec) def submit_load( self, job_id: int, src_spec: LoadStoreSpec, - dst_spec: GPULoadStoreSpec, + device_ptrs: DevicePointers | GPULoadStoreSpec, expedited: bool = False, ) -> bool: # vLLM main submits step-synchronous loads with `expedited=True`; this # worker keeps a single load path and honours the same contract. - return self._load_handler.transfer_async(job_id, src_spec, dst_spec) + return self._load_handler.transfer_async(job_id, src_spec, device_ptrs) def get_finished(self) -> list[TransferResult]: return self._store_handler.get_finished() + self._load_handler.get_finished() diff --git a/vllm_ascend/distributed/weight_transfer/hccl_engine.py b/vllm_ascend/distributed/weight_transfer/hccl_engine.py index ca71afa58ebf..e77b51192e22 100644 --- a/vllm_ascend/distributed/weight_transfer/hccl_engine.py +++ b/vllm_ascend/distributed/weight_transfer/hccl_engine.py @@ -9,8 +9,9 @@ """ from concurrent.futures import ThreadPoolExecutor -from contextlib import suppress +from contextlib import nullcontext, suppress from dataclasses import dataclass +from importlib import import_module from typing import TYPE_CHECKING, Any, ClassVar import torch @@ -154,13 +155,16 @@ def receive_weights(self, update_info: HCCLWeightTransferUpdateInfo) -> None: # else: update_info = self._last_hccl_update_info.get(self.model, update_info) - from vllm.model_executor.model_loader.mtp_validation import ( - disable_mtp_completeness_check, - ) + mtp_check = nullcontext() + if vllm_version_is("0.33.0"): + # The pinned main tree used by mypy no longer contains this + # release-only module. Resolve it only in the release runtime. + mtp_validation = import_module("vllm.model_executor.model_loader.mtp_validation") + mtp_check = mtp_validation.disable_mtp_completeness_check() # The worker is handed an explicit NPU device, so bind it before touching # any stream: the packed helpers resolve the ambient device. - with torch.npu.device(self.device), disable_mtp_completeness_check(): + with torch.npu.device(self.device), mtp_check: metadata = zip(update_info.names, update_info.shapes, update_info.dtype_names) if params["packed"]: packed_broadcast_consumer( diff --git a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py index 3832d824ecf4..762d0a98c5f2 100644 --- a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py +++ b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py @@ -4,8 +4,10 @@ import os import socket +from contextlib import nullcontext from dataclasses import dataclass from functools import lru_cache +from importlib import import_module from typing import TYPE_CHECKING, Any, ClassVar import torch @@ -26,6 +28,7 @@ packed_npu_ipc_consumer, packed_npu_ipc_producer, ) +from vllm_ascend.utils import vllm_version_is if TYPE_CHECKING: from vllm.config import VllmConfig @@ -198,11 +201,14 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: weight = rebuild_npu_tensor(*list_args) weights.append((name, weight)) - from vllm.model_executor.model_loader.mtp_validation import ( - disable_mtp_completeness_check, - ) + mtp_check = nullcontext() + if vllm_version_is("0.33.0"): + # The pinned main tree used by mypy no longer contains this + # release-only module. Resolve it only in the release runtime. + mtp_validation = import_module("vllm.model_executor.model_loader.mtp_validation") + mtp_check = mtp_validation.disable_mtp_completeness_check() - with disable_mtp_completeness_check(): + with mtp_check: self.model.load_weights(weights) diff --git a/vllm_ascend/models/deepseek_v4/vl_model.py b/vllm_ascend/models/deepseek_v4/vl_model.py index d30a8e289d3f..06b20ca8d16b 100644 --- a/vllm_ascend/models/deepseek_v4/vl_model.py +++ b/vllm_ascend/models/deepseek_v4/vl_model.py @@ -34,6 +34,7 @@ DeepseekV4Aligner, DeepseekV4ViT, ) +from vllm_ascend.utils import vllm_version_is def _vision_parameter_name(name: str) -> str | None: @@ -288,6 +289,9 @@ def language_weights() -> Iterator[tuple[str, torch.Tensor]]: return loaded_vision | {f"language_model.{name}" for name in loaded_language} def process_weights_after_loading(self) -> None: + # #57678 already finalizes language_model in the children-first walk. + if not vllm_version_is("0.33.0"): + return hook = getattr( self.language_model, "process_weights_after_loading", diff --git a/vllm_ascend/models/deepseek_v41/engram/hash_state.py b/vllm_ascend/models/deepseek_v41/engram/hash_state.py index ad989a676381..c0b935cf2a42 100644 --- a/vllm_ascend/models/deepseek_v41/engram/hash_state.py +++ b/vllm_ascend/models/deepseek_v41/engram/hash_state.py @@ -23,10 +23,14 @@ DEAD_ID, EngramLayout, NgramHashState, - _write_hash_cache_kernel, ) from vllm.triton_utils import triton +from vllm_ascend.utils import vllm_version_is + +if vllm_version_is("0.33.0"): + from vllm.models.deepseek_v41.common.engram import _write_hash_cache_kernel + class AscendEngramSlotCache: """``swa_cache_module`` view for upstream ``NgramHashState``.""" @@ -76,6 +80,15 @@ class AscendNgramHashState(NgramHashState): Keep that compiler workaround in this subclass, without changing vLLM. """ + if not vllm_version_is("0.33.0"): + # Main removed the MRv1 history cache; Ascend callers still share the + # slotless and legacy launch adapter while v0.33.0 is supported. + use_slot_cache = False + + @property + def block_size(self) -> int: + return self.swa_cache_module.block_size + def forward( self, input_ids: torch.Tensor, @@ -84,8 +97,8 @@ def forward( dead_mask: torch.Tensor, lookback_token_ids: torch.Tensor, lookback_dead_mask: torch.Tensor, - slot_mapping: torch.Tensor | None, - block_table: torch.Tensor | None, + slot_mapping: torch.Tensor | None = None, + block_table: torch.Tensor | None = None, ) -> torch.Tensor: """Compute [tokens, layers, hash columns] int32 n-gram hashes. diff --git a/vllm_ascend/models/deepseek_v41/vl_model.py b/vllm_ascend/models/deepseek_v41/vl_model.py index ade0a8d099a6..36c2efc8187e 100644 --- a/vllm_ascend/models/deepseek_v41/vl_model.py +++ b/vllm_ascend/models/deepseek_v41/vl_model.py @@ -24,6 +24,8 @@ ) from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm_ascend.utils import vllm_version_is + from .model import AscendDeepseekV41LLMForCausalLM from .vision import DeepseekV41Aligner, DeepseekV41ViT @@ -323,6 +325,9 @@ def language_weights() -> Iterator[tuple[str, torch.Tensor]]: return loaded_vision | {f"language_model.{name}" for name in loaded_language} def process_weights_after_loading(self) -> None: + # #57678 already finalizes language_model in the children-first walk. + if not vllm_version_is("0.33.0"): + return hook = getattr( self.language_model, "process_weights_after_loading", diff --git a/vllm_ascend/models/gemma4_mm.py b/vllm_ascend/models/gemma4_mm.py index 8b175127fdc2..85f5aa5bae1d 100644 --- a/vllm_ascend/models/gemma4_mm.py +++ b/vllm_ascend/models/gemma4_mm.py @@ -18,6 +18,7 @@ import torch from transformers import AutoModel from vllm.config import VllmConfig +from vllm.model_executor.models import gemma4_mm from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM from vllm.model_executor.models.gemma4_mm import ( Gemma4ForConditionalGeneration, @@ -30,6 +31,7 @@ AscendW4A4MXFP4DynamicLinearMethod, AscendW8A8MXFP8DynamicLinearMethod, ) +from vllm_ascend.utils import vllm_version_is # NOTE: The upstream Gemma4 vision patch embedder may cast pixel_values to the # projection weight dtype. For Ascend MXFP4/MXFP8 dynamic quantization, the @@ -85,9 +87,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # Initialize the attribute here to prevent AttributeError when processing image inputs. lora_config = vllm_config.lora_config self._enable_mm_lora = bool(lora_config is not None and lora_config.enable_tower_connector_lora) + self.use_data_parallel = bool( + multimodal_config is not None and getattr(multimodal_config, "mm_encoder_tp_mode", None) == "data" + ) with self._mark_tower_model(vllm_config, {"image", "video"}): self.vision_tower = AutoModel.from_config(config=config.vision_config) + # v0.33.0 predates the device-side normalization added in #65577. + if not vllm_version_is("0.33.0"): + self.vision_input_norm = gemma4_mm.build_mm_input_norm(vllm_config.model_config) self.embed_vision = Gemma4MultimodalEmbedder( config.vision_config, config.text_config, diff --git a/vllm_ascend/models/glm5next/kda.py b/vllm_ascend/models/glm5next/kda.py index 3de260a3d93e..3a1e3bad192d 100644 --- a/vllm_ascend/models/glm5next/kda.py +++ b/vllm_ascend/models/glm5next/kda.py @@ -195,6 +195,9 @@ def __init__( # MergedColumnParallelLinear's constructor hardcodes the TP group. # Retain its shard loaders while allocating through the PCP column layer. PCPColumnParallelLinear.__init__(self, input_size, sum(self.output_sizes), **kwargs) + # vLLM #63752: the PCP allocation bypasses the stacked initializer. + self.shard_ids = tuple(range(len(self.output_sizes))) + self.loaded_shard_ids: dict[str, set[str | int]] = {} class Glm5NextLinearAttention(GatedDeltaNetAttention): diff --git a/vllm_ascend/models/kimi_k3/kimi_k3.py b/vllm_ascend/models/kimi_k3/kimi_k3.py index 6d22959a8dc4..9141298ae101 100644 --- a/vllm_ascend/models/kimi_k3/kimi_k3.py +++ b/vllm_ascend/models/kimi_k3/kimi_k3.py @@ -100,7 +100,7 @@ from vllm_ascend.quantization.methods.base import PreparedLinearInput from vllm_ascend.quantization.methods.w8a8.w8a8_dynamic import AscendW8A8DynamicLinearMethod from vllm_ascend.quantization.methods.w8a8.w8a8_mxfp8 import AscendW8A8MXFP8DynamicLinearMethod -from vllm_ascend.utils import get_rotation_path, is_950 +from vllm_ascend.utils import get_rotation_path, is_950, vllm_version_is from vllm_ascend.worker.v2.pp_transport import ( PPTransportDataType, add_pp_transport_tensors, @@ -1047,9 +1047,24 @@ def _cache_aux_pp_layout(self) -> None: # A materialized state at start_layer belongs to this stage, # unlike the raw state emitted by the preceding stage. self._aux_slot_base_cached = bisect_left(self.aux_hidden_state_layers, self.start_layer) - if pp.is_last_rank: + if vllm_version_is("0.33.0") and pp.is_last_rank: self._aux_upstream_total_cached = self._aux_slot_base_cached + if not vllm_version_is("0.33.0"): + + def aux_hidden_state_slot_keys(self) -> tuple[tuple[str, ...], tuple[str, ...]]: + if self.config.attn_res_block_size is None or not self.dspark_aux_capture_materialized: + return super().aux_hidden_state_slot_keys() + pp = get_pp_group() + if pp.world_size < 2: + return (), () + # Materialized DSpark captures before layer L. A cut at L belongs + # to the receiving stage, unlike raw post-layer capture. + base = self._aux_slot_base_cached + end = base if pp.is_last_rank else bisect_left(self.aux_hidden_state_layers, self.end_layer) + key = self.AUX_HIDDEN_STATE_KEY + return tuple(f"{key}{i}" for i in range(base)), tuple(f"{key}{i}" for i in range(base, end)) + def make_empty_intermediate_tensors( self, batch_size: int, @@ -1146,7 +1161,13 @@ def forward( residual = intermediate_tensors["residual"] materialized_aux = self.dspark_aux_capture_materialized - aux_hidden_states = get_pp_transport_tensors(intermediate_tensors, PPTransportDataType.AUX_HIDDEN_STATES) + # Main's runner relays upstream slots and merges them once on the last + # stage. The release model still carries the cumulative payload. + aux_hidden_states = ( + get_pp_transport_tensors(intermediate_tensors, PPTransportDataType.AUX_HIDDEN_STATES) + if vllm_version_is("0.33.0") + else [] + ) if self.use_sequence_parallel: if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): forward_context = get_forward_context() @@ -1229,8 +1250,11 @@ def forward( hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] residual = sp_all_gather(residual)[:full_num_tokens] aux_hidden_states = [sp_all_gather(aux)[:full_num_tokens] for aux in aux_hidden_states] + output = IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + if not vllm_version_is("0.33.0"): + return IntermediateTensors({**output.tensors, **self.pack_local_aux_hidden_states(aux_hidden_states)}) return add_pp_transport_tensors( - IntermediateTensors({"hidden_states": hidden_states, "residual": residual}), + output, PPTransportDataType.AUX_HIDDEN_STATES, aux_hidden_states, ) diff --git a/vllm_ascend/models/layer/attention/layer.py b/vllm_ascend/models/layer/attention/layer.py index 76876ee17334..2924e394799d 100644 --- a/vllm_ascend/models/layer/attention/layer.py +++ b/vllm_ascend/models/layer/attention/layer.py @@ -187,6 +187,7 @@ def __init__( ] self.kv_cache_dtype = kv_cache_dtype + self.act_dtype = torch.get_default_dtype() self.use_sparse = True def forward( @@ -198,7 +199,10 @@ def forward( ) -> torch.Tensor: return q - def process_weights_after_loading(self, act_dtype: torch.dtype): + def process_weights_after_loading(self, act_dtype: torch.dtype | None = None): + # #57678 calls attention hooks without arguments on main. + if act_dtype is None: + act_dtype = self.act_dtype if hasattr(self.impl, "process_weights_after_loading"): self.impl.process_weights_after_loading(act_dtype) diff --git a/vllm_ascend/models/minimax_m3/minimax_m3.py b/vllm_ascend/models/minimax_m3/minimax_m3.py index ff5b2bc949fc..7d4760ac79b8 100644 --- a/vllm_ascend/models/minimax_m3/minimax_m3.py +++ b/vllm_ascend/models/minimax_m3/minimax_m3.py @@ -116,7 +116,7 @@ _register_m3_sparse_packed_modules, _use_fused_qkv_indexer, ) -from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type +from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type, vllm_version_is from vllm_ascend.worker.v2.pp_transport import ( PPTransportDataType, add_pp_transport_tensors, @@ -1181,9 +1181,10 @@ def forward( assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] - aux_hidden_states = get_pp_transport_tensors( - intermediate_tensors, - PPTransportDataType.AUX_HIDDEN_STATES, + aux_hidden_states = ( + get_pp_transport_tensors(intermediate_tensors, PPTransportDataType.AUX_HIDDEN_STATES) + if vllm_version_is("0.33.0") + else [] ) hidden_states_are_sp = False @@ -1231,6 +1232,10 @@ def forward( hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]] residual = sp_all_gather(residual)[: positions.shape[0]] intermediate_tensors = IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + if not vllm_version_is("0.33.0"): + return IntermediateTensors( + {**intermediate_tensors.tensors, **self.pack_local_aux_hidden_states(aux_hidden_states)} + ) return add_pp_transport_tensors( intermediate_tensors, PPTransportDataType.AUX_HIDDEN_STATES, diff --git a/vllm_ascend/models/minimax_m3/msa_m3.py b/vllm_ascend/models/minimax_m3/msa_m3.py index 19ae51416e11..7c11d36f93cb 100644 --- a/vllm_ascend/models/minimax_m3/msa_m3.py +++ b/vllm_ascend/models/minimax_m3/msa_m3.py @@ -63,6 +63,7 @@ ) from vllm_ascend.ops.linear import AscendColumnParallelLinear from vllm_ascend.ops.linear_op import get_parallel_op +from vllm_ascend.utils import vllm_version_is # Minimum used by heterogeneous-KV alignment; the final FP8 cache block is 128. _MINIMAX_M3_FP8_KV_CACHE_BLOCK_SIZE = 64 @@ -1149,6 +1150,8 @@ def update_graph_params( class AscendMiniMaxM3QKVParallelLinearWithIndexer(QKVParallelLinear): """Fused [q | k | v | index_q | index_k] column-parallel GEMM for M3 sparse layers.""" + shard_ids = ("q", "k", "v", "index_q", "index_k") + def __init__( self, hidden_size: int, @@ -1181,6 +1184,8 @@ def __init__( self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) self.num_kv_head_replicas = 1 self.num_index_heads = self.num_kv_heads + # #63752: this class bypasses StackedColumnParallelLinear.__init__. + self.loaded_shard_ids: dict[str, set[str | int]] = {} q = self.num_heads * self.head_size kv = self.num_kv_heads * self.head_size @@ -1210,15 +1215,18 @@ def forward(self, input_): return self.custom_op.apply(input_) return super().forward(input_) - def validate_shard_id(self, loaded_shard_id: str | None) -> None: - if loaded_shard_id is None: - return - if loaded_shard_id not in ("q", "k", "v", "index_q", "index_k"): + def validate_shard_id(self, shard_id: Any, param: Parameter | None = None) -> bool: + if not vllm_version_is("0.33.0"): + return super().validate_shard_id(shard_id, param) + if shard_id is None: + return True + if shard_id not in self.shard_ids: raise ValueError( "Shard id for AscendMiniMaxM3QKVParallelLinearWithIndexer must be " "one of 'q', 'k', 'v', 'index_q', 'index_k'; got " - f"{loaded_shard_id}." + f"{shard_id}." ) + return True def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None: h = self.head_size @@ -1247,7 +1255,7 @@ def weight_loader_v2( loaded_weight, loaded_shard_id: str | None = None, ) -> None: - self.validate_shard_id(loaded_shard_id) + self.validate_shard_id(loaded_shard_id, param) assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k") shard_offset = self._get_shard_offset_mapping(loaded_shard_id) @@ -1273,7 +1281,7 @@ def weight_loader( loaded_weight, loaded_shard_id: str | None = None, ) -> None: - self.validate_shard_id(loaded_shard_id) + self.validate_shard_id(loaded_shard_id, param) assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k") output_dim = getattr(param, "output_dim", None) assert output_dim is not None diff --git a/vllm_ascend/ops/fused_moe/fused_moe.py b/vllm_ascend/ops/fused_moe/fused_moe.py index 52dcfe7bd327..e16aa330f2b1 100644 --- a/vllm_ascend/ops/fused_moe/fused_moe.py +++ b/vllm_ascend/ops/fused_moe/fused_moe.py @@ -41,6 +41,7 @@ SharedExpertParallelMode, SubmittedSharedExpertGateUp, ) +from vllm_ascend.utils import vllm_version_is def _ascend_moe_forward_shared_sp_fake( @@ -96,7 +97,18 @@ def __init__( routed_input_transform=None, routed_output_transform=None, routed_scaling_factor=1, + divide_shared_output_in_fp16=False, + routed_scale_mode=None, ): + scaling_kwargs = {} + if not vllm_version_is("0.33.0"): + from vllm.model_executor.layers.fused_moe.config import RoutedScaleMode + + if routed_scale_mode is None: + routed_scale_mode = ( + RoutedScaleMode.OUTPUT_FP16_GUARD if divide_shared_output_in_fp16 else RoutedScaleMode.ROUTER + ) + scaling_kwargs["routed_scale_mode"] = routed_scale_mode super().__init__( layer_name, moe_config, @@ -109,6 +121,7 @@ def __init__( routed_input_transform, routed_output_transform, routed_scaling_factor, + **scaling_kwargs, ) self._gate = gate self.hidden_size = moe_config.hidden_dim diff --git a/vllm_ascend/ops/fused_moe/routed_experts.py b/vllm_ascend/ops/fused_moe/routed_experts.py index 6aa1f714d1fc..88b28e56b724 100644 --- a/vllm_ascend/ops/fused_moe/routed_experts.py +++ b/vllm_ascend/ops/fused_moe/routed_experts.py @@ -362,6 +362,12 @@ def _empty_like_preserving_npu_format(tensor: torch.Tensor) -> torch.Tensor: class EplbExpertTensorList(list[torch.Tensor]): """Per-expert tensors exposed through the upstream EPLB weight contract.""" + @property + def is_cuda(self) -> bool: + # NPU expert lists take the native per-row copy path rather than the + # CUDA kernel, which requires a single contiguous weight allocation. + return self[0].is_cuda + @property def shape(self) -> torch.Size: return torch.Size((len(self), *self[0].shape)) diff --git a/vllm_ascend/ops/kda_parallel.py b/vllm_ascend/ops/kda_parallel.py index 9a28a431d305..0a15ab8576e3 100644 --- a/vllm_ascend/ops/kda_parallel.py +++ b/vllm_ascend/ops/kda_parallel.py @@ -58,6 +58,9 @@ def __init__(self, input_size: int, output_sizes: list[int], **kwargs) -> None: self.output_sizes = output_sizes.copy() # Keep the merged checkpoint loaders, but allocate with explicit ranks. KDAColumnParallelLinear.__init__(self, input_size, sum(self.output_sizes), **kwargs) + # vLLM #63752: custom allocation bypasses the stacked initializer. + self.shard_ids = tuple(range(len(self.output_sizes))) + self.loaded_shard_ids: dict[str, set[str | int]] = {} class KDAInputParallelLinear(_KimiGDNMergedColumnParallelLinear, KDAMergedColumnParallelLinear): diff --git a/vllm_ascend/ops/linear.py b/vllm_ascend/ops/linear.py index 625965909aaa..d8939631ac35 100644 --- a/vllm_ascend/ops/linear.py +++ b/vllm_ascend/ops/linear.py @@ -238,6 +238,8 @@ def __init__( return_bias=return_bias, disable_tp=disable_tp, ) + # vLLM #63752: custom allocation bypasses the stacked initializer. + self.loaded_shard_ids: dict[str, set[str | int]] = {} def forward( self, @@ -278,6 +280,7 @@ def __init__( # TODO(realliujiaxu): Replace the initialization code below with super().__init__ after # linear of vllm supports custom comm group self.output_sizes = output_sizes + self.shard_ids = tuple(range(len(output_sizes))) assert all(output_size % self.tp_size == 0 for output_size in output_sizes) AscendColumnParallelLinear.__init__( self, @@ -292,6 +295,8 @@ def __init__( return_bias=return_bias, disable_tp=disable_tp, ) + # vLLM #63752: custom allocation bypasses the stacked initializer. + self.loaded_shard_ids: dict[str, set[str | int]] = {} def forward( self, diff --git a/vllm_ascend/ops/mla.py b/vllm_ascend/ops/mla.py index 9676a8208e03..96b6684592ab 100644 --- a/vllm_ascend/ops/mla.py +++ b/vllm_ascend/ops/mla.py @@ -36,6 +36,7 @@ from vllm_ascend.attention.indexer import AscendSFAIndexerBackend from vllm_ascend.attention.mla_v1 import AscendMLAAttention from vllm_ascend.attention.utils import mark_fused_preprocess_weights +from vllm_ascend.utils import vllm_version_is class IndexerWrapper(nn.Module): @@ -190,6 +191,7 @@ def __init__( # both layernorms as extra args, so this flag has no effect here and is # accepted for signature compatibility only. fuse_qkv_rmsnorm: bool = False, + shardq_eligible: bool = False, ) -> None: nn.Module.__init__(self) self.hidden_size = hidden_size @@ -207,7 +209,9 @@ def __init__( self.skip_topk = skip_topk # This is an upstream CUDA indexer hint. Ascend accepts it to preserve # constructor compatibility, but its indexer does not consume it. - del allow_short_prefill_indexer_scoring_skip + # ShardQ eligibility is another CUDA implementation hint (#63956). + # Ascend retains its own attention/indexer layout for both values. + del allow_short_prefill_indexer_scoring_skip, shardq_eligible hf_config = get_current_vllm_config().model_config.hf_text_config self.tp_size = get_tensor_model_parallel_world_size() self.layers = hf_config.num_hidden_layers @@ -260,14 +264,18 @@ def __init__( original_process_weights = self.mla_attn.process_weights_after_loading - def wrapped_process_weights(act_dtype: torch.dtype): + def wrapped_process_weights(act_dtype: torch.dtype | None = None): from vllm_ascend.attention.sfa_v1 import AscendSFAImpl + if act_dtype is None: + act_dtype = self.mla_attn.act_dtype if self.dcp_q_replicate_enabled: self.mla_attn.impl.process_weights_after_loading(act_dtype) elif not isinstance(self.mla_attn.impl, AscendSFAImpl): - # Both supported vLLM versions dispatch to the impl here. - original_process_weights(act_dtype) + if vllm_version_is("0.33.0"): + original_process_weights(act_dtype) + else: + original_process_weights() else: # SFA disposes kv_b_proj, so bypass upstream's dense packing. self.mla_attn.impl.process_weights_after_loading(act_dtype) diff --git a/vllm_ascend/ops/triton/mamba/postprocess.py b/vllm_ascend/ops/triton/mamba/postprocess.py index e1e55af2f1fe..35c7ddf73790 100644 --- a/vllm_ascend/ops/triton/mamba/postprocess.py +++ b/vllm_ascend/ops/triton/mamba/postprocess.py @@ -95,6 +95,9 @@ def postprocess_mamba_fused_kernel( # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds # the post-step new_num_computed value (V2 supplies the advanced count). PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False, + # #55688 delegates ReplaySSM temporal migration to its ring materializer. + # Convolution snapshots and accepted-count normalization remain generic. + SKIP_TEMPORAL_STATE_COPY: tl.constexpr = False, # TEMPORAL_TILES: when > 1, the temporal copy body is partitioned across # TEMPORAL_TILES CTAs along the u64 inner range. Callers must launch a # 3D grid (num_reqs, total_states, TEMPORAL_TILES). Default 1 preserves @@ -222,6 +225,10 @@ def postprocess_mamba_fused_kernel( if src_block_idx == dest_block_idx and accept_token_bias == 0: return + if SKIP_TEMPORAL_STATE_COPY: + if conv_width == 0: + return + offsets = tl.arange(0, COPY_BLOCK_SIZE) if CONV_STATE_DIM_FIRST and is_conv_state: dim_row_count = tl.load(state_dim_row_count_ptr + state_idx) diff --git a/vllm_ascend/ops/triton/mamba/precopy.py b/vllm_ascend/ops/triton/mamba/precopy.py index b1bf7d15c484..4da1d2fcbf6b 100644 --- a/vllm_ascend/ops/triton/mamba/precopy.py +++ b/vllm_ascend/ops/triton/mamba/precopy.py @@ -131,6 +131,7 @@ def precopy_mamba_align_fused_kernel( CONV_STATE_DIM_FIRST: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, TEMPORAL_TILES: tl.constexpr = 1, + ZERO_TEMPORAL_TOKEN_BIAS: tl.constexpr = False, ): batch_idx = tl.program_id(0) state_idx = tl.program_id(1) @@ -154,6 +155,9 @@ def precopy_mamba_align_fused_kernel( return token_bias = tl.load(token_bias_ptr + req_idx) + if ZERO_TEMPORAL_TOKEN_BIAS: + conv_width = tl.load(state_conv_widths_ptr + state_idx) + token_bias = tl.where(conv_width > 0, token_bias, 0) _copy_mamba_state_block( state_idx, diff --git a/vllm_ascend/ops/triton/sample/categorical_sample.py b/vllm_ascend/ops/triton/sample/categorical_sample.py index 673d2e07b54f..ee9f0f6991c8 100644 --- a/vllm_ascend/ops/triton/sample/categorical_sample.py +++ b/vllm_ascend/ops/triton/sample/categorical_sample.py @@ -424,6 +424,7 @@ def _categorical_sample_main( use_fp64: bool = False, logits_cache_source: torch.Tensor | None = None, logits_cache_row_mapping: torch.Tensor | None = None, + bias: torch.Tensor | None = None, ) -> torch.Tensor: # 0.33.0 has `logits_cache_source` (only watermarking supplies a separate # source); main adds the per-row mapping for batch-sharded caches. The NPU @@ -432,6 +433,11 @@ def _categorical_sample_main( raise NotImplementedError("Batch-sharded draft logits are not supported on NPU.") if logits_cache_source is not None: raise NotImplementedError("A separate draft-logits cache source is not supported on NPU.") + if bias is not None: + # DSpark main (#64476) may defer its Markov bias to the sampler. Match + # the upstream logits-dtype rounding before temperature, sampling and + # draft-cache writes; keep the caller's logits unchanged. + logits = (logits.float() + bias.float()).to(logits.dtype) return _categorical_sample( logits, expanded_idx_mapping, diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 20b08d240f9d..d8f1a2c33680 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -174,6 +174,10 @@ # available in the supported vLLM version. Retain only the NPU backend, # operator, communicator, load normalization, and weight views. # +# Main2main November 12: vLLM #66859 replaces routed-scale booleans with +# RoutedScaleMode. Translate Ascend model callers before router construction +# and forward native modes to both runners. Remove the boolean bridge when +# the matched release also uses the enum. # ** 8. File: platform/patch_fused_moe.py** # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 1. `vllm.model_executor.layers.fused_moe.FusedMoEFactory` @@ -894,6 +898,21 @@ # Remove only this embedding provisioning when the minimum supported vLLM # version provides it in DeepseekV2Model; keep the unrelated PP transport. # +# 3. `vllm.model_executor.models.deepseek_v2.DeepseekV2Model` +# Why: +# vLLM main #66177 identifies aux decoders by EagleModelMixin and moves +# remote-state aggregation into PPHandler. Copying release-only fields +# fails at import and does not provide the required decoder identity. +# How: +# Keep the existing member copy on v0.33.0. On main, replace the module's +# decoder class with an EagleModelMixin subclass of the decorated, +# Ascend-patched native decoder. Native slot bookkeeping and PPHandler +# own local/remote relay; Ascend retains capture and TP gathering. +# Related PR: +# https://github.com/vllm-project/vllm/pull/66177 +# Future Plan: +# Remove the replacement when native DeepseekV2Model inherits the mixin. +# # ** 3. File: worker/patch_distributed.py** # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 1. `vllm.distributed.parallel_state.GroupCoordinator` @@ -1512,8 +1531,8 @@ # ** 25. File: worker/patch_eagle_speculator.py** # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 1. `vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator.SpeculatorCudaGraphManager` -# (named `vllm.v1.worker.gpu.spec_decode.autoregressive.speculator` before -# upstream vLLM #60335; both supported lanes use the new path). +# 2. `vllm.v1.worker.gpu.spec_decode.autoregressive.speculator.SpeculatorCudaGraphManager` +# (vLLM main #60989 constructs the decode manager in the shared AR base). # Why: # The shared v2 autoregressive speculative-decoding path (Eagle/MTP) uses # an upstream CUDA graph manager, which is not usable on Ascend. @@ -1638,6 +1657,8 @@ # before device filtering. Preserve the original dictionary and container # types for model bindings and connector registration; remove the old # first-tensor wrapper in patch_attn_utils.py. +# Preserve #55688's ReplaySSM tracker collection for inherited scheduler +# block copies when the supported upstream exports that collector. # Related PR (if no, explain why): # No upstream PR linked; this adapts Ascend multi-tensor allocations. # Future Plan: @@ -1771,3 +1792,36 @@ # How: initialize the capability-gated per-layer NPU adaptation after upstream init. # Related PR: vllm-project/vllm#60050. # Future Plan: remove this initialization hook when a shared backend extension supports it. +# ** File: platform/patch_kimi_k31_config.py ** +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# 1. `vllm.transformers_utils.config._CONFIG_REGISTRY["kimi_k31"]` +# Why: +# Expose the K3.1 config identity and Ascend DSA constraints. Native main +# #66030 consumes a KimiLinearDSAConfig, whereas v0.33.0 uses our dict. +# How: +# Register the existing platform config subclass once. Keep the explicit +# release dict branch; delegate main's dataclass construction, layer +# lookup and serialization to native KimiLinearConfig. +# Related PR: +# https://github.com/vllm-project/vllm/pull/66030 +# Future Plan: +# Remove this registration/subclass when native kimi_k31 registration +# provides the required config identity and Ascend DSA validation. +# +# ** File: worker/patch_spec_decode_310.py ** +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# `vllm_ascend.worker.v2.spec_decode.init_speculator` +# Why: +# Route regular MTP to the existing 310P implementation, while keeping +# Gemma4 and Step3.5 on their existing fallback paths. Native #63373 +# removes use_step3p5_mtp from main's SpeculativeConfig with MRv1. +# How: +# Share uses_step3p5_mtp with the main Ascend router: explicit v0.33.0 +# uses the native predicate; main reads the same draft HF model_type. +# No native method is reinstalled and constructor behavior is unchanged. +# Related PR: +# https://github.com/vllm-project/vllm/pull/63373 +# Future Plan: +# Remove the compatibility predicate when supported native versions +# expose a common model-routing capability or share one MTP route. +# diff --git a/vllm_ascend/patch/platform/patch_balance_schedule.py b/vllm_ascend/patch/platform/patch_balance_schedule.py index 3b7af0b69a85..b6cc6538ccd4 100644 --- a/vllm_ascend/patch/platform/patch_balance_schedule.py +++ b/vllm_ascend/patch/platform/patch_balance_schedule.py @@ -78,6 +78,8 @@ from vllm.v1.utils import record_function_or_nullcontext from vllm_ascend.ascend_config import init_ascend_config +from vllm_ascend.core.scheduler_output import make_cached_request_data +from vllm_ascend.utils import vllm_version_is def _balance_scheduling_enabled(vllm_config) -> bool: @@ -762,16 +764,8 @@ def _skip_request(from_queue) -> None: # len(self.running). assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len(scheduled_running_reqs) <= len(self.running) - # Get the longest common prefix among all requests in the running queue. - # This can be potentially used for cascade attention. - num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) - with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): - if self.running: - any_request_id = self.running[0].request_id - num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) - # Construct the scheduler output. - if self.use_v2_model_runner: + if not vllm_version_is("0.33.0") or self.use_v2_model_runner: scheduled_new_reqs.extend(scheduled_resumed_reqs) scheduled_resumed_reqs.clear() new_reqs_data = [ @@ -794,16 +788,20 @@ def _skip_request(from_queue) -> None: ] with record_function_or_nullcontext("schedule: make_cached_request_data"): - cached_reqs_data = self._make_cached_request_data( + cached_reqs_data, schedule_metadata = make_cached_request_data( + self, scheduled_running_reqs, scheduled_resumed_reqs, num_scheduled_tokens, scheduled_spec_decode_tokens, req_to_new_blocks, + new_reqs_data, + scheduled_encoder_inputs, + encoder_compute_budget, ) # Record the request ids that were scheduled in this step (MRV1-only). - if not self.use_v2_model_runner: + if vllm_version_is("0.33.0") and not self.use_v2_model_runner: self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) @@ -839,7 +837,7 @@ def _skip_request(from_queue) -> None: total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, - num_common_prefix_blocks=num_common_prefix_blocks, + **schedule_metadata, preempted_req_ids={req.request_id for req in preempted_reqs}, # finished_req_ids is an existing state in the scheduler, # instead of being newly scheduled in this step. diff --git a/vllm_ascend/patch/platform/patch_fused_moe.py b/vllm_ascend/patch/platform/patch_fused_moe.py index 032d293109a7..bc9b3a98021f 100644 --- a/vllm_ascend/patch/platform/patch_fused_moe.py +++ b/vllm_ascend/patch/platform/patch_fused_moe.py @@ -43,6 +43,7 @@ from vllm_ascend.distributed.eplb.eplb_state import AscendEplbLayerState from vllm_ascend.ops.fused_moe.eplb import use_allgather_eplb from vllm_ascend.ops.fused_moe.router.router_factory import create_ascend_fused_moe_router +from vllm_ascend.utils import vllm_version_is _EPLB_ROUTER_ADAPTED = "_vllm_ascend_eplb_router_adapted" @@ -124,6 +125,8 @@ def _ascend_FusedMoE( num_redundant_experts: int = 0, n_shared_experts: int | None = None, apply_routed_scale_to_output: bool = False, + divide_shared_output_in_fp16: bool = False, + routed_scale_mode: Any | None = None, hash_indices_table: torch.Tensor | None = None, runner_cls: Any | None = None, runner_args: dict[str, Any] | None = None, @@ -134,6 +137,25 @@ def _ascend_FusedMoE( image_sentinel_lo: int = 129257, **kwargs, ): + # Ascend-owned model callers retain the release booleans; native main + # callers pass the enum introduced by vLLM #66859. Normalize before + # constructing either the router or the native factory to scale once. + scaling_kwargs: dict[str, Any] + if vllm_version_is("0.33.0"): + scaling_kwargs = {"apply_routed_scale_to_output": apply_routed_scale_to_output} + else: + from vllm.model_executor.layers.fused_moe.config import RoutedScaleMode + + if routed_scale_mode is None: + routed_scale_mode = ( + RoutedScaleMode.OUTPUT_FP16_GUARD + if apply_routed_scale_to_output and divide_shared_output_in_fp16 + else RoutedScaleMode.OUTPUT + if apply_routed_scale_to_output + else RoutedScaleMode.ROUTER + ) + apply_routed_scale_to_output = routed_scale_mode.scales_output + scaling_kwargs = {"routed_scale_mode": routed_scale_mode} # RoutedExperts allocates its parameters before AscendMoERunner is # constructed. Propagate Ascend EPLB capacity into the upstream factory so # redundant expert slots are present when weights are created and loaded. @@ -185,7 +207,7 @@ def _ascend_FusedMoE( enable_eplb=enable_eplb, num_redundant_experts=num_redundant_experts, n_shared_experts=n_shared_experts, - apply_routed_scale_to_output=apply_routed_scale_to_output, + **scaling_kwargs, hash_indices_table=hash_indices_table, runner_cls=runner_cls, runner_args=runner_args, diff --git a/vllm_ascend/patch/platform/patch_kimi_k31_config.py b/vllm_ascend/patch/platform/patch_kimi_k31_config.py index 975d2cf916f2..05320251b65a 100644 --- a/vllm_ascend/patch/platform/patch_kimi_k31_config.py +++ b/vllm_ascend/patch/platform/patch_kimi_k31_config.py @@ -6,14 +6,21 @@ from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm_ascend.utils import vllm_version_is + class KimiK31TextConfig(KimiLinearConfig): # Preserve K3.1 identity when with_hf_config enters the kimi_linear model. is_kimi_k31 = True def __init__(self, dsa_config: dict | None = None, **kwargs): - super().__init__(**kwargs) - self.dsa_config = dsa_config + if vllm_version_is("0.33.0"): + super().__init__(**kwargs) + self.dsa_config = dsa_config + else: + # Native main owns the dataclass consumed by the composite config. + super().__init__(dsa_config=dsa_config, **kwargs) + dsa_config = super().to_dict().get("dsa_config") if dsa_config is None: return self.index_topk = dsa_config["index_topk"] @@ -36,6 +43,8 @@ def __init__(self, dsa_config: dict | None = None, **kwargs): raise ValueError("Kimi K3.1 KDA and DSA layers must be disjoint") def is_dsa_layer(self, layer_idx: int) -> bool: + if not vllm_version_is("0.33.0"): + return super().is_dsa_layer(layer_idx) return self.dsa_config is not None and layer_idx + 1 in self.dsa_config["dsa_layers"] diff --git a/vllm_ascend/patch/platform/patch_pp_mtp.py b/vllm_ascend/patch/platform/patch_pp_mtp.py index 222481ee9f28..d9e26ca87430 100644 --- a/vllm_ascend/patch/platform/patch_pp_mtp.py +++ b/vllm_ascend/patch/platform/patch_pp_mtp.py @@ -344,7 +344,8 @@ def _apply_patch() -> None: _patch_model_runner_output() # The patch for engine core has been moved to patch_engine_core.py _patch_scheduler_update_after_schedule() - _patch_scheduler_make_cached_request_data() + if vllm_version_is("0.33.0"): + _patch_scheduler_make_cached_request_data() _patch_scheduler_update_from_output() _patch_model_config_validation() diff --git a/vllm_ascend/patch/platform/patch_speculative_config.py b/vllm_ascend/patch/platform/patch_speculative_config.py index 7fb7cad3eb99..814d654a8d51 100644 --- a/vllm_ascend/patch/platform/patch_speculative_config.py +++ b/vllm_ascend/patch/platform/patch_speculative_config.py @@ -93,10 +93,11 @@ def _normalize_legacy_qwen3_dspark_config( return _orig_hf_config_override( # type: ignore[call-arg] hf_config, transformers_mtp_overrides ) + # #63205 renames this legacy architecture before returning to our hook. + legacy_qwen3_dspark = hf_config.model_type == "qwen3" and "DSparkDraftModel" in (hf_config.architectures or ()) hf_config = _orig_hf_config_override(hf_config) _normalize_kimi_dflash_rope(hf_config) - architectures = hf_config.architectures or () - if hf_config.model_type == "qwen3" and "DSparkDraftModel" in architectures: + if legacy_qwen3_dspark: dflash_config = hf_config.dflash_config hf_config.update( { diff --git a/vllm_ascend/patch/platform/patch_swa_bounded_replay.py b/vllm_ascend/patch/platform/patch_swa_bounded_replay.py index f0579c4082dc..d7e44e9105ca 100644 --- a/vllm_ascend/patch/platform/patch_swa_bounded_replay.py +++ b/vllm_ascend/patch/platform/patch_swa_bounded_replay.py @@ -53,6 +53,8 @@ from vllm.v1.core.sched.output import CachedRequestData from vllm.v1.core.sched.scheduler import Scheduler +from vllm_ascend.utils import vllm_version_is + _EXPECTED_PARAMETERS = ( "self", "running_reqs", @@ -150,7 +152,7 @@ def _make_cached_request_data( Scheduler._make_cached_request_data = _make_cached_request_data -if _upstream_carries_replay_start(): +if vllm_version_is("0.33.0") and _upstream_carries_replay_start(): # vLLM loads general plugins in the engine-core process and in every worker # (``load_general_plugins``), so this module reaches both sides of the # msgpack boundary -- the encoder needs the field declared to write it, the diff --git a/vllm_ascend/patch/worker/__init__.py b/vllm_ascend/patch/worker/__init__.py index f8af8c3b08bc..8b9a2876b3d6 100644 --- a/vllm_ascend/patch/worker/__init__.py +++ b/vllm_ascend/patch/worker/__init__.py @@ -87,6 +87,10 @@ ) _ar_spec.SpeculatorCudaGraphManager = AutoRegressiveAclGraphManager310 + if not vllm_version_is("0.33.0"): + from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as _shared_ar_spec + + _shared_ar_spec.SpeculatorCudaGraphManager = AutoRegressiveAclGraphManager310 # only patch routed experts capture in main2main. import vllm_ascend.patch.worker.patch_routed_experts_capture # noqa diff --git a/vllm_ascend/patch/worker/patch_bind_kv_cache.py b/vllm_ascend/patch/worker/patch_bind_kv_cache.py index 73a697f11c3b..d2574fb56ad9 100644 --- a/vllm_ascend/patch/worker/patch_bind_kv_cache.py +++ b/vllm_ascend/patch/worker/patch_bind_kv_cache.py @@ -7,6 +7,7 @@ from vllm.v1.worker.utils import defaultdict, extract_layer_index from vllm_ascend.attention.dsa_v41 import DeepseekV41CacheLayer +from vllm_ascend.utils import vllm_version_is # Without this patch, it will raise an exception when initialize kv_cache. @@ -17,6 +18,7 @@ def bind_kv_cache( runner_kv_caches: list[torch.Tensor], num_attn_module: int = 1, kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, + layer_kv_caches: dict[str, torch.Tensor] | None = None, ) -> None: """ Bind the allocated KV cache to both ModelRunner and forward context so @@ -51,10 +53,12 @@ def bind_kv_cache( ordered_layer_names.append(layer_name) # Bind kv_caches to forward context - for layer_name, kv_cache in kv_caches.items(): + for layer_name, kv_cache in (layer_kv_caches or kv_caches).items(): forward_context[layer_name].kv_cache = kv_cache - # vLLM #52506 adds ReplaySSM ring trackers on main. - utils.share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups) + # #55688 moves ReplaySSM tracker ownership out of cache binding. Older + # vLLM versions still require the legacy sharing hook after binding. + if vllm_version_is("0.33.0"): + utils.share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups) def bind_kv_cache_to_layers( @@ -78,4 +82,5 @@ def bind_kv_cache_to_layers( else: layer.kv_cache = kv_cache ordered_layer_names = sorted(kv_caches, key=lambda name: extract_layer_index(name, num_attn_module)) - utils.share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups) + if vllm_version_is("0.33.0"): + utils.share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups) diff --git a/vllm_ascend/patch/worker/patch_deepseek_v2.py b/vllm_ascend/patch/worker/patch_deepseek_v2.py index a8ecaa0a0574..2c28ac7c22db 100644 --- a/vllm_ascend/patch/worker/patch_deepseek_v2.py +++ b/vllm_ascend/patch/worker/patch_deepseek_v2.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.models import deepseek_v2 as deepseek_v2_module from vllm.model_executor.models.deepseek_v2 import ( DeepseekV2ForCausalLM, DeepSeekV2FusedQkvAProjLinear, @@ -39,7 +40,7 @@ from vllm_ascend.attention.context_parallel.common_cp import use_dcp_q_replicate from vllm_ascend.ops.linear import AscendDCPGroupColumnParallelLinear -from vllm_ascend.utils import is_mtp_layer, should_reuse_topk +from vllm_ascend.utils import is_mtp_layer, should_reuse_topk, vllm_version_is from vllm_ascend.worker.v2 import pp_transport from vllm_ascend.worker.v2.pp_transport import ( PPTransportDataType, @@ -340,17 +341,33 @@ def _patched_deepseek_v2_model_init(self, *args, **kwargs): # Reuse upstream slot bookkeeping; Ascend owns capture and TP gathering. -for _member in ( - "AUX_HIDDEN_STATE_KEY", - "_aux_slot_base_cached", - "_aux_upstream_total_cached", - "_set_aux_hidden_state_layers", - "_cache_aux_pp_layout", - "pack_local_aux_hidden_states", - "collect_remote_aux_hidden_states", -): - setattr(DeepseekV2Model, _member, getattr(EagleModelMixin, _member)) -DeepseekV2Model.supports_aux_hidden_states_over_pp = True +if vllm_version_is("0.33.0"): + for _member in ( + "AUX_HIDDEN_STATE_KEY", + "_aux_slot_base_cached", + "_aux_upstream_total_cached", + "_set_aux_hidden_state_layers", + "_cache_aux_pp_layout", + "pack_local_aux_hidden_states", + "collect_remote_aux_hidden_states", + ): + setattr(DeepseekV2Model, _member, getattr(EagleModelMixin, _member)) + DeepseekV2Model.supports_aux_hidden_states_over_pp = True +else: + + class AscendEagleDeepseekV2Model(DeepseekV2Model, EagleModelMixin): + """Expose the native decoder identity required by vLLM #66177. + + Inherit the decorated DeepSeek implementation and its Ascend patches; + the main mixin supplies local slots, and the runner merges remote ones. + Remove this replacement when native DeepSeek inherits EagleModelMixin. + """ + + supports_aux_hidden_states_over_pp = True + + deepseek_v2_module.DeepseekV2Model = AscendEagleDeepseekV2Model + # Native constructors use this cached factory, not the module alias above. + DeepseekV2ForCausalLM.model_cls = AscendEagleDeepseekV2Model def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: diff --git a/vllm_ascend/patch/worker/patch_eagle_speculator.py b/vllm_ascend/patch/worker/patch_eagle_speculator.py index 6137688adabe..096e766ed68d 100644 --- a/vllm_ascend/patch/worker/patch_eagle_speculator.py +++ b/vllm_ascend/patch/worker/patch_eagle_speculator.py @@ -16,8 +16,17 @@ # This file is a part of the vllm-ascend project. # -from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import speculator as vllm_speculator_module +from vllm_ascend.utils import vllm_version_is from vllm_ascend.worker.v2.spec_decode.target_dependent_ar.aclgraph import AutoRegressiveAclGraphManager -vllm_speculator_module.SpeculatorCudaGraphManager = AutoRegressiveAclGraphManager +if vllm_version_is("0.33.0"): + from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import speculator as ar_speculator +else: + from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as ar_speculator + from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import speculator as target_speculator + + # #60989 moved decode graph construction to the shared AR base class. + target_speculator.SpeculatorCudaGraphManager = AutoRegressiveAclGraphManager + +ar_speculator.SpeculatorCudaGraphManager = AutoRegressiveAclGraphManager diff --git a/vllm_ascend/patch/worker/patch_mamba_utils.py b/vllm_ascend/patch/worker/patch_mamba_utils.py index 5364dbd6e997..5583b38ca431 100644 --- a/vllm_ascend/patch/worker/patch_mamba_utils.py +++ b/vllm_ascend/patch/worker/patch_mamba_utils.py @@ -2,11 +2,11 @@ import torch from vllm.v1.worker import mamba_utils -from vllm.v1.worker.mamba_utils import MambaCopyBuffers from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile from vllm_ascend.ops.triton.batch_memcpy import batch_memcpy_kernel from vllm_ascend.ops.triton.mamba.postprocess import postprocess_mamba_fused_kernel +from vllm_ascend.utils import vllm_version_is # Upstream uses 16 temporal-copy tiles to saturate H100/GB200. K3 already # exposes 138 independent state programs per request, while Triton-Ascend @@ -47,17 +47,18 @@ def _batch_memcpy_triton(src_ptrs, dst_ptrs, sizes): # Ascend NPU does not support DT_UINT64 in aclnnInplaceZero. # MambaCopyBuffers.create() uses torch.uint64 for src_ptrs/dst_ptrs, # which triggers a runtime error. Remap to int64 at the source. -_original_create = MambaCopyBuffers.create +if vllm_version_is("0.33.0"): + from vllm.v1.worker.mamba_utils import MambaCopyBuffers + _original_create = MambaCopyBuffers.create -@classmethod -def _patched_create(cls, max_num_reqs, kv_cache_config, copy_funcs, make_buffer): - return _original_create( - max_num_reqs, - kv_cache_config, - copy_funcs, - lambda n, dtype: make_buffer(n, dtype=torch.int64 if dtype == torch.uint64 else dtype), - ) + @classmethod + def _patched_create(cls, max_num_reqs, kv_cache_config, copy_funcs, make_buffer): + return _original_create( + max_num_reqs, + kv_cache_config, + copy_funcs, + lambda n, dtype: make_buffer(n, dtype=torch.int64 if dtype == torch.uint64 else dtype), + ) - -MambaCopyBuffers.create = _patched_create + MambaCopyBuffers.create = _patched_create diff --git a/vllm_ascend/patch/worker/patch_model_runner.py b/vllm_ascend/patch/worker/patch_model_runner.py index 830d644c5d4c..7ca04aeaaae5 100644 --- a/vllm_ascend/patch/worker/patch_model_runner.py +++ b/vllm_ascend/patch/worker/patch_model_runner.py @@ -142,9 +142,13 @@ def initialize_kv_cache( self.kv_cache_config, use_replayssm=self.cache_config.use_replayssm, ) - piecewise_capture_available = bool( - upstream.envs.VLLM_USE_BREAKABLE_CUDAGRAPH or upstream.has_compiled_submodule(self.model) - ) + if vllm_version_is("0.33.0"): + breakable_cg = upstream.envs.VLLM_USE_BREAKABLE_CUDAGRAPH + graph_kwargs = {"decode_query_len": self.decode_query_len} + else: + breakable_cg = upstream.is_breakable_cudagraph_enabled(self.compilation_config) + graph_kwargs = {"decode_query_lens": self.vllm_config.get_decode_cudagraph_query_lens(self.decode_query_len)} + piecewise_capture_available = bool(breakable_cg or upstream.has_compiled_submodule(self.model)) if self.adaptive_verification is not None: self.compilation_config.cudagraph_mode = upstream.resolve_adaptive_cudagraph_mode( self.compilation_config.cudagraph_mode, @@ -155,7 +159,7 @@ def initialize_kv_cache( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, self.decode_query_len, - use_v2_model_runner=True, + **({"use_v2_model_runner": True} if vllm_version_is("0.33.0") else {}), kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, @@ -166,7 +170,7 @@ def initialize_kv_cache( self.vllm_config, self.device, cudagraph_mode, - decode_query_len=self.decode_query_len, + **graph_kwargs, lora_capture_cases=self.lora_capture_cases, varlen_decode=self.adaptive_verification is not None, ubatch_runner=self.ubatch_runner, @@ -208,6 +212,11 @@ def initialize_kv_cache( for tensor in (cache if isinstance(cache, (tuple, list)) else (cache,)) if tensor is not None and tensor.device == self.device ] + # #55688 adds the collector on main; v0.33.0 shares trackers during binding. + if not vllm_version_is("0.33.0"): + self.replayssm_block_copy_tensors = upstream.get_replayssm_block_copy_tensors( + self.compilation_config.static_forward_context + ) if is_profiling: self.kv_connector = upstream.NO_OP_KV_CONNECTOR else: diff --git a/vllm_ascend/patch/worker/patch_spec_decode_310.py b/vllm_ascend/patch/worker/patch_spec_decode_310.py index 9d8650d69a7a..c61b7dd4d5b7 100644 --- a/vllm_ascend/patch/worker/patch_spec_decode_310.py +++ b/vllm_ascend/patch/worker/patch_spec_decode_310.py @@ -31,7 +31,7 @@ def _init_speculator_310p(vllm_config, device, req_states=None): if ( speculative_config.method == "mtp" and not speculative_config.use_gemma4_mtp() - and not speculative_config.use_step3p5_mtp() + and not ascend_spec_decode.uses_step3p5_mtp(speculative_config) ): from vllm_ascend._310p.worker.v2.spec_decode.mtp_speculator import ( AscendMTPSpeculator310, diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index 2f6e901699b9..a164d25b74a8 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -65,6 +65,7 @@ refresh_block_size, update_cudagraph_capture_sizes, enable_sp, + vllm_version_is, ) if TYPE_CHECKING: @@ -1430,6 +1431,18 @@ def _setup_compile_backend( set(compilation_config.cudagraph_capture_sizes or []) | {reduced_cg_cap} ) compilation_config.max_cudagraph_capture_size = None + # Main resolves separate prefill/decode sizes on the first pass. They + # were derived from the truncated grid, so regenerate them alongside + # the restored off-stride default rather than mixing two size sets. + # Explicit split sizes never receive reduced_cg_cap (see below). + if not vllm_version_is("0.33.0"): + for name in ( + "prefill_cudagraph_capture_token_sizes", + "max_prefill_cudagraph_capture_tokens", + "decode_cudagraph_capture_request_sizes", + "max_decode_cudagraph_capture_requests", + ): + setattr(compilation_config, name, None) delattr(compilation_config, "reduced_cg_cap") # Recompute cudagraph sizes before extending splitting_ops (honors the @@ -1444,7 +1457,13 @@ def _setup_compile_backend( and not vllm_config.model_config.enforce_eager and requires_tp_aligned_capture_sizes ): - original_sizes = compilation_config.cudagraph_capture_sizes + # Main V2 resolves prefill tokens and decode requests separately. + # The TP token-layout restriction applies to the prefill grid; keep + # the resolved decode request grid used by the graph managers intact. + prefill_sizes = None + if not vllm_version_is("0.33.0"): + prefill_sizes = compilation_config.prefill_cudagraph_capture_token_sizes + original_sizes = prefill_sizes if prefill_sizes is not None else compilation_config.cudagraph_capture_sizes tp_size = vllm_config.parallel_config.tensor_parallel_size sp_aclgraph_sizes = [size for size in original_sizes if size % tp_size == 0] if not sp_aclgraph_sizes: @@ -1454,7 +1473,17 @@ def _setup_compile_backend( f"{vllm_config.parallel_config.tensor_parallel_size}" ) - if len(sp_aclgraph_sizes) != len(original_sizes): + if prefill_sizes is not None: + compilation_config.prefill_cudagraph_capture_token_sizes = sp_aclgraph_sizes + compilation_config.max_prefill_cudagraph_capture_tokens = sp_aclgraph_sizes[-1] + # Re-record the token union from the resolved split grids. Keeping + # a legacy-only list (e.g. [4] with one decode request) makes the + # engine handshake's config re-entry reject the derived sizes. + vllm_config._split_cudagraph_capture_sizes( + False, vllm_config.get_decode_cudagraph_query_lens(vllm_config.uniform_decode_query_len) + ) + compilation_config.post_init_cudagraph_sizes() + elif len(sp_aclgraph_sizes) != len(original_sizes): # Match max_cudagraph_capture_size with the valid max size to avoid # initialization error of vllm server. compilation_config.max_cudagraph_capture_size = sp_aclgraph_sizes[-1] @@ -1561,9 +1590,40 @@ def _setup_worker_and_scheduler( "vllm_ascend.core.batch_job_aware_scheduler.BatchJobAwareScheduler" ) + _validate_copied_scheduler_async_kv_limit(vllm_config, ascend_config) _downgrade_swa_bounded_replay(vllm_config, ascend_config) +def _copied_scheduler_mode(ascend_config) -> str | None: + """Name a mode that replaces upstream's schedule implementation.""" + scheduler_config = ascend_config.scheduler_config + if scheduler_config.recompute_scheduler_enable: + return "recompute_scheduler_enable" + if scheduler_config.dyntra_lb_config.enabled: + return "dyntra_lb_config" + if scheduler_config.profiling_chunk_config.enabled: + return "profiling_chunk_config" + if scheduler_config.enable_balance_scheduling: + return "enable_balance_scheduling" + return None + + +def _validate_copied_scheduler_async_kv_limit(vllm_config: VllmConfig, ascend_config) -> None: + """Reject the new async-queue opt-in until copied schedulers drain it.""" + if vllm_version_is("0.33.0"): + return + if getattr(vllm_config.scheduler_config, "max_async_kv_loads_per_step", None) is None: + return + mode = _copied_scheduler_mode(ascend_config) + if mode is not None: + # vLLM #61725 routes these requests into async_kv_waiting. The copied + # schedulers only scan waiting/kv_holding_waiting, leaving them stranded. + raise ValueError( + f"max_async_kv_loads_per_step is not supported with {mode}. " + "Leave this limit unset or disable the Ascend scheduler extension." + ) + + def _swa_bounded_replay_unsupported_mode(ascend_config) -> str | None: """Name the Ascend mode that cannot honor SWA bounded replay. @@ -1580,16 +1640,7 @@ def _swa_bounded_replay_unsupported_mode(ascend_config) -> str | None: its ``super().schedule()`` fallback only applies while balancing is off. """ - scheduler_config = ascend_config.scheduler_config - if scheduler_config.recompute_scheduler_enable: - return "recompute_scheduler_enable" - if scheduler_config.dyntra_lb_config.enabled: - return "dyntra_lb_config" - if scheduler_config.profiling_chunk_config.enabled: - return "profiling_chunk_config" - if scheduler_config.enable_balance_scheduling: - return "enable_balance_scheduling" - return None + return _copied_scheduler_mode(ascend_config) def _downgrade_swa_bounded_replay(vllm_config: VllmConfig, ascend_config) -> None: @@ -1838,6 +1889,17 @@ def _get_reduced_cg_cap(vllm_config: VllmConfig) -> int | None: return None if compilation_config.cudagraph_capture_sizes is not None: return None + if not vllm_version_is("0.33.0") and any( + getattr(compilation_config, name) is not None + for name in ( + "prefill_cudagraph_capture_token_sizes", + "max_prefill_cudagraph_capture_tokens", + "decode_cudagraph_capture_request_sizes", + "max_decode_cudagraph_capture_requests", + "max_decode_cudagraph_capture_tokens", + ) + ): + return None scheduler_config = getattr(vllm_config, "scheduler_config", None) max_num_seqs = getattr(scheduler_config, "max_num_seqs", None) diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index c8aeffa26e37..430e00cf5c37 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -736,6 +736,21 @@ def setup_ascend_local_comm_res(user_device_id: int, kv_transfer_config: Any | N os.environ["ASCEND_LOCAL_COMM_RES"] = json.dumps(data, ensure_ascii=False, separators=(",", ":")) +def get_config_hash(config: Any, *, cached: bool = False) -> str: + """Use the selected vLLM configuration hash contract. + + Mutable copies, such as normalized KV cache configs, must stay uncached. + Remove the compute_hash branch when v0.33.0 support ends. + """ + if vllm_version_is("0.33.0"): + return config.compute_hash() + # utils is imported while the out-of-tree platform is initializing. + # Importing vllm.config there re-enters current_platform resolution. + from vllm.config import utils as config_utils + + return config_utils.hash_config_cached(config) if cached else config_utils.hash_config(config) + + @functools.cache @torch._dynamo.disable def vllm_version_is(target_vllm_version: str): diff --git a/vllm_ascend/worker/v2/aclgraph_utils.py b/vllm_ascend/worker/v2/aclgraph_utils.py index bac086e19f94..c75e5c48383c 100644 --- a/vllm_ascend/worker/v2/aclgraph_utils.py +++ b/vllm_ascend/worker/v2/aclgraph_utils.py @@ -16,7 +16,7 @@ # limitations under the License. # This file is a part of the vllm-ascend project. # -from collections.abc import Callable +from collections.abc import Callable, Mapping from contextlib import contextmanager from typing import Any @@ -46,7 +46,7 @@ ContextSource, UpdatableGraph, ) -from vllm_ascend.utils import use_updatable_graph +from vllm_ascend.utils import use_updatable_graph, vllm_version_is from vllm_ascend.worker.v2.utils import communicator_switch @@ -88,7 +88,7 @@ def __init__( vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, - decode_query_len: int, + decode_query_len: int | Mapping[int, int], model_runner: Any, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, @@ -215,6 +215,7 @@ def capture( """Capture CUDA graphs for model forward pass.""" metadata = getattr(model_state, "device_metadata", None) model = ModelWithContext(model, device_metadata=metadata) + legacy_kwargs = {"has_lora": has_lora} if vllm_version_is("0.33.0") else {} try: with communicator_switch(): return super().capture( @@ -226,7 +227,7 @@ def capture( attn_groups, kv_cache_config, pcp_manager=pcp_manager, - has_lora=has_lora, + **legacy_kwargs, use_aux_hidden_state_outputs=use_aux_hidden_state_outputs, lora_capture_hook=lora_capture_hook, progress_bar_desc=progress_bar_desc, diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index d55cfd9cbe11..2c2f1342f5b8 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -1607,6 +1607,8 @@ def allocate_kv_cache_main( device: torch.device, layout: Any, kernel_block_sizes: list[int], + *, + allocate: Callable[[int], torch.Tensor] | None = None, ) -> dict[str, Any]: """Allocate Ascend KV cache through vLLM main's #51718 entry point. @@ -1615,6 +1617,8 @@ def allocate_kv_cache_main( still consumes separate K/V (and backend-specific state) tensors, so keep the Ascend allocation/reshape contract behind the new entry point. """ + # #56492 always forwards this keyword; NPU disables extensible caches. + assert allocate is None, "Extensible KV cache allocation is not supported on Ascend" del layout vllm_config = get_current_vllm_config() attn_layers = get_layers_from_vllm_config(vllm_config, AttentionLayerBase) diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index b6375c477c51..4e676c1b1f9c 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -19,6 +19,7 @@ import inspect import logging +from collections.abc import Mapping from contextlib import AbstractContextManager, contextmanager, nullcontext from copy import deepcopy from typing import Any @@ -160,6 +161,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # boundary and sequence length padding during FULL graph execution. self.use_fia = False self._uses_flashmla = False + # #56492: inherited dummy runs inspect this even on non-VMM platforms. + self.extensible_kv_cache = None # FusedMoE can be constructed by the parent initializer and reads this # capacity while setting up MC2 communication. set_potential_max_tokens(vllm_config) @@ -196,7 +199,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE and ( self.compilation_config.mode == CompilationMode.VLLM_COMPILE - or breakable_cudagraph.is_breakable_cudagraph_enabled() + or ( + breakable_cudagraph.is_breakable_cudagraph_enabled() + if vllm_version_is("0.33.0") + else breakable_cudagraph.is_breakable_cudagraph_enabled(self.compilation_config) + ) ) and not self.model_config.enforce_eager ) @@ -1591,16 +1598,20 @@ def factory( # type: ignore[misc] vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, - decode_query_len: int, + decode_query_len: int | None = None, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, - ubatch_runner: Any = None, # vLLM main (#51700) + ubatch_runner: Any = None, + *, + decode_query_lens: Mapping[int, int] | None = None, ): + query_shapes = decode_query_len if vllm_version_is("0.33.0") else decode_query_lens + assert query_shapes is not None return ModelAclGraphManager( vllm_config, device, cudagraph_mode, - decode_query_len, + query_shapes, model_runner, lora_capture_cases=lora_capture_cases, varlen_decode=varlen_decode, diff --git a/vllm_ascend/worker/v2/model_states/mamba_hybrid.py b/vllm_ascend/worker/v2/model_states/mamba_hybrid.py index f8ff2af6fbff..79c227e4003c 100644 --- a/vllm_ascend/worker/v2/model_states/mamba_hybrid.py +++ b/vllm_ascend/worker/v2/model_states/mamba_hybrid.py @@ -40,7 +40,7 @@ ) from vllm.v1.worker.utils import AttentionGroup -from vllm_ascend.utils import model_uses_kimi_kda, model_uses_kpool_indexer +from vllm_ascend.utils import model_uses_kimi_kda, model_uses_kpool_indexer, vllm_version_is from vllm_ascend.worker.v2.attn_utils import build_attn_metadata from vllm_ascend.worker.v2.input_batch import AscendInputBatch from vllm_ascend.worker.v2.model_states.default import AscendModelState @@ -222,12 +222,19 @@ def _ensure_align_ctx( copy_funcs = self.model.get_mamba_state_copy_funcs(mamba_types) validate_mamba_state_copy_funcs(mamba_groups, copy_funcs) self._mamba_state_copy_funcs = copy_funcs + # vLLM #55688 counts state tensors from each layer's bindings. + # v0.33 still uses the constructor without forward_context. + context_kwargs: dict[str, Any] = {} + if vllm_version_is("0.33.0"): + context_kwargs["make_buffer"] = lambda n, dtype: CpuGpuBuffer(n, dtype=dtype, device=self.device) + else: + context_kwargs["forward_context"] = self.vllm_config.compilation_config.static_forward_context self._mamba_ctx = MambaSpecDecodeGPUContext.create( max_num_reqs=self.max_num_reqs, kv_cache_config=kv_cache_config, copy_funcs=self._mamba_state_copy_funcs, device=self.device, - make_buffer=lambda n, dtype: CpuGpuBuffer(n, dtype=dtype, device=self.device), + **context_kwargs, ) ctx = self._mamba_ctx if not ctx.is_initialized: diff --git a/vllm_ascend/worker/v2/spec_decode/__init__.py b/vllm_ascend/worker/v2/spec_decode/__init__.py index 98b3978dc810..f68a2da2daf1 100644 --- a/vllm_ascend/worker/v2/spec_decode/__init__.py +++ b/vllm_ascend/worker/v2/spec_decode/__init__.py @@ -22,10 +22,24 @@ import torch from vllm.config import CompilationMode, VllmConfig +from vllm_ascend.utils import vllm_version_is + if TYPE_CHECKING: from vllm.v1.worker.gpu.states import RequestState +def uses_step3p5_mtp(speculative_config) -> bool: + if vllm_version_is("0.33.0"): + return speculative_config.use_step3p5_mtp() + # Native #63373 removed this predicate with MRv1. Retain the same model + # distinction so Step3.5 keeps its existing fallback instead of generic MTP. + return ( + speculative_config.method == "mtp" + and speculative_config.draft_model_config is not None + and getattr(speculative_config.draft_model_config.hf_config, "model_type", None) == "step3p5_mtp" + ) + + def init_speculator( vllm_config: VllmConfig, device: torch.device, @@ -89,7 +103,7 @@ def init_speculator( ) return AscendGemma4Speculator(vllm_config, device) - if speculative_config.method == "mtp" and not speculative_config.use_step3p5_mtp(): + if speculative_config.method == "mtp" and not uses_step3p5_mtp(speculative_config): from vllm_ascend.worker.v2.spec_decode.mtp.speculator import ( AscendMTPSpeculator, ) diff --git a/vllm_ascend/worker/v2/spec_decode/dflash/aclgraph.py b/vllm_ascend/worker/v2/spec_decode/dflash/aclgraph.py index 170ed293f765..21a18983a3e7 100644 --- a/vllm_ascend/worker/v2/spec_decode/dflash/aclgraph.py +++ b/vllm_ascend/worker/v2/spec_decode/dflash/aclgraph.py @@ -24,7 +24,7 @@ ContextSource, UpdatableGraph, ) -from vllm_ascend.utils import use_updatable_graph +from vllm_ascend.utils import use_updatable_graph, vllm_version_is from vllm_ascend.worker.v2.aclgraph_utils import collect_sorted_captured_token_sizes, model_capture_wrapper from vllm_ascend.worker.v2.utils import communicator_switch @@ -35,14 +35,17 @@ def __init__( vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, - decode_query_len: int, + decode_query_len: int | Mapping[int, int] | None = None, speculator: Any = None, + *, + decode_query_lens: Mapping[int, int] | None = None, ): + query_shapes = decode_query_len if vllm_version_is("0.33.0") else decode_query_lens super().__init__( vllm_config, device, cudagraph_mode, - decode_query_len, + query_shapes, ) # It is set by AscendDFlashSpeculator.init_cudagraph_manager after creation, diff --git a/vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py b/vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py index 8a4117d42f9e..a329ea583bb7 100644 --- a/vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py +++ b/vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py @@ -346,6 +346,9 @@ def rejection_sample( posterior_threshold: float = 0.95, posterior_alpha: float = 0.4, draft_logits_num_shards: int = 1, + # vLLM #65100 hints that masked vocab blocks may be skipped. The NPU + # resampler keeps its full residual scan, which preserves the same result. + skip_masked_vocab_blocks: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # Upstream passes logits after applying sampling parameters. Resolve the # worker configuration here to keep its sampler and chunking flow unchanged. diff --git a/vllm_ascend/worker/v2/spec_decode/standalone_ar/speculator.py b/vllm_ascend/worker/v2/spec_decode/standalone_ar/speculator.py index b54138945391..217eafa84f81 100644 --- a/vllm_ascend/worker/v2/spec_decode/standalone_ar/speculator.py +++ b/vllm_ascend/worker/v2/spec_decode/standalone_ar/speculator.py @@ -37,12 +37,10 @@ def _build_attn_metadata( slot_mappings: torch.Tensor | None = None, ) -> dict[str, Any] | None: draft_seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound - if step == 0: - # TODO: Remove this workaround once upstream includes the correction - # token in the standalone AR prefill sequence-length bound. - # Draft prefill appends the target's correction token to the - # expanded query, so the upstream builder must receive target - # length + 1. Keep the input batch's bounds unchanged. + if step == 0 and vllm_version_is("0.33.0"): + # #60989 includes the correction token before calling this builder + # on main. Only release still needs the extra prefill token here. + # Keep the input batch's bounds unchanged. draft_seq_lens_cpu_upper_bound = (seq_lens_cpu_upper_bound[:num_reqs] + 1).clamp(max=self.max_model_len) num_tokens = int(query_start_loc_np[-1]) diff --git a/vllm_ascend/worker/v2/spec_decode/target_dependent_ar/aclgraph.py b/vllm_ascend/worker/v2/spec_decode/target_dependent_ar/aclgraph.py index 53e797c39811..ee8b702627c1 100644 --- a/vllm_ascend/worker/v2/spec_decode/target_dependent_ar/aclgraph.py +++ b/vllm_ascend/worker/v2/spec_decode/target_dependent_ar/aclgraph.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any import torch @@ -17,12 +17,9 @@ ) from vllm.v1.worker.gpu.input_batch import InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState -from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( - SpeculatorCudaGraphManager, -) from vllm.v1.worker.utils import AttentionGroup -from vllm_ascend.ascend_forward_context import _EXTRA_CTX +from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType, select_moe_comm_method from vllm_ascend.compilation.acl_graph import ( set_draft_graph_params, set_draft_graph_prefill_params, @@ -32,13 +29,20 @@ SharedSource, UpdatableGraph, ) -from vllm_ascend.utils import use_updatable_graph +from vllm_ascend.utils import is_drafter_moe_model, use_updatable_graph, vllm_version_is from vllm_ascend.worker.v2.aclgraph_utils import ( collect_sorted_captured_token_sizes, model_capture_wrapper, ) from vllm_ascend.worker.v2.utils import communicator_switch +if vllm_version_is("0.33.0"): + from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( # type: ignore[import-not-found] + SpeculatorCudaGraphManager, + ) +else: + from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import SpeculatorCudaGraphManager + if TYPE_CHECKING: from vllm.v1.worker.gpu.pcp_manager import PCPManager @@ -64,21 +68,53 @@ def __init__( vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, - decode_query_len: int, + decode_query_len: int | Mapping[int, int] | None = None, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + *, + decode_query_lens: Mapping[int, int] | None = None, + decode_request_sizes: Sequence[int] | None = None, ): # Upstream uses num_speculative_steps + 1 as the draft-prefill query # length and 1 for draft decode. + # vLLM #63406 keys uniform decode graphs by request count. + graph_kwargs = {} + if vllm_version_is("0.33.0"): + query_shapes = decode_query_len + else: + if decode_query_lens is None: + assert isinstance(decode_query_len, Mapping) + decode_query_lens = decode_query_len + query_shapes = decode_query_lens + decode_query_len = max(decode_query_lens) + if ( + decode_query_len == 1 + and decode_request_sizes is not None + and cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + and is_drafter_moe_model(vllm_config) + ): + # Main expands this independent draft grid beyond the target + # capture bound (#63406). ALLTOALL preprocess synchronizes + # dynamic token counts on the host and cannot run in FULL. + # Native dispatch keeps larger batches eager, with all their + # requests/tokens, when they have no compatible capture size. + decode_request_sizes = [ + size + for size in decode_request_sizes + if select_moe_comm_method(size, vllm_config) != MoECommType.ALLTOALL + ] + graph_kwargs["decode_request_sizes"] = decode_request_sizes + assert isinstance(decode_query_len, int) if decode_query_len > 1: cudagraph_mode = get_draft_prefill_cudagraph_mode(vllm_config, cudagraph_mode) super().__init__( vllm_config, device, cudagraph_mode, - decode_query_len, + query_shapes, lora_capture_cases=lora_capture_cases, varlen_decode=varlen_decode, + **graph_kwargs, ) # Upstream constructs graph managers without a speculator reference. diff --git a/vllm_ascend/worker/v2/spec_decode/target_dependent_ar/speculator.py b/vllm_ascend/worker/v2/spec_decode/target_dependent_ar/speculator.py index 2852a35666e9..04ae4d33c056 100644 --- a/vllm_ascend/worker/v2/spec_decode/target_dependent_ar/speculator.py +++ b/vllm_ascend/worker/v2/spec_decode/target_dependent_ar/speculator.py @@ -47,7 +47,7 @@ from vllm_ascend.attention.indexer import AscendSFAIndexerBackend from vllm_ascend.attention.mla_v1 import AscendMLABackend, AscendMLAMetadata from vllm_ascend.attention.sfa_v1 import AscendSFABackend -from vllm_ascend.utils import vllm_version_is +from vllm_ascend.utils import get_config_hash, vllm_version_is from vllm_ascend.worker.dcp_utils import DCPManager from vllm_ascend.worker.v2.aclgraph_utils import _get_graph_update_backend from vllm_ascend.worker.v2.attn_utils import ( @@ -223,7 +223,7 @@ def set_eplb_state(self, eplb_state: Any) -> None: # Only graph-external phase hooks update draft source validity. if self._eplb_source_is_padding is None: self._eplb_source_is_padding = torch.ones_like(self.input_buffers.is_padding) - model_state = eplb_state.model_states[self.speculative_config.draft_model_config.compute_hash()] + model_state = eplb_state.model_states[get_config_hash(self.speculative_config.draft_model_config, cached=True)] for layer in model_state.model.moe_layers: layer.eplb_state.allgather_source_is_padding = self._eplb_source_is_padding @@ -275,7 +275,7 @@ def propose( # shape. Negotiate its DP agreement instead of reusing the target's. sync_state = None if self.pcp_manager is not None else dp_sync_state if self.pcp_manager is not None: - attn_metadata = self._get_pcp_draft_prefill_metadata(attn_metadata) + attn_metadata = self._get_draft_prefill_metadata(attn_metadata) # wrap build_attn_metadata to use Ascend attention metadata building. # so we can call super().propose() directly. with ( @@ -307,10 +307,10 @@ def propose( ) @staticmethod - def _get_pcp_draft_prefill_metadata(attn_metadata: dict[str, Any]) -> dict[str, Any]: + def _get_draft_prefill_metadata(attn_metadata: dict[str, Any]) -> dict[str, Any]: """Keep eager draft MLA queries within the target's real decode rows. - FULL target graphs include query offsets for padded PCP requests. The + FULL target graphs include query offsets for padded requests. The draft prefill runs outside FULL graphs and MLA slices its query to num_decode_tokens, so padding requests must have zero query length. Copy only the changed metadata: target graph updates still need the @@ -433,18 +433,32 @@ def _generate_draft( FIA operator.""" # vLLM #60335 added `num_speculative_steps` on main. with self._lmhead_pad_size_context(capture_num_reqs=num_reqs): - super()._generate_draft( - num_reqs, - num_tokens_padded, - attn_metadata, - slot_mappings, - num_tokens_across_dp, - cudagraph_runtime_mode, - num_speculative_steps, - ) + if vllm_version_is("0.33.0"): + super()._generate_draft( + num_reqs, + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + num_speculative_steps, + ) + else: + super()._decode( + num_reqs, + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + num_speculative_steps, + ) if attn_metadata is not None: self._update_decode_attn_metadata(attn_metadata, 1, num_reqs) + # #60989 invokes _decode through the new shared AR loop. + _decode = _generate_draft + def _multi_step_decode( self, num_reqs: int, @@ -494,6 +508,12 @@ def _prefill( name: metadata for name, metadata in attn_metadata.items() if name in self.draft_attn_layer_names } + # #63406 can introduce request-padded target graphs even without PCP. + # Eager/PW MLA slices queries to real decode rows; its cumulative query + # lengths must exclude target padding too. Keep FULL graph metadata. + if not vllm_version_is("0.33.0") and attn_metadata is not None and cudagraph_runtime_mode != CUDAGraphMode.FULL: + attn_metadata = self._get_draft_prefill_metadata(attn_metadata) + with self._lmhead_pad_size_context(capture_num_reqs=num_reqs): # vLLM main forwards `skip_sample` so the drafter KV stays in sync # without sampling when num_speculative_tokens == 0. diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py index 83ae8b4866e2..57a2f63ad258 100644 --- a/vllm_ascend/worker/worker.py +++ b/vllm_ascend/worker/worker.py @@ -118,9 +118,11 @@ register_ascend_fla_customop, register_device_print, setup_ascend_local_comm_res, + vllm_version_is, ) +# Both supported versions resolve grouping through the platform's builder. def get_kv_cache_groups( vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec], @@ -430,7 +432,10 @@ def finish_weight_update(self) -> None: assert self.weight_transfer_engine is not None self.weight_transfer_engine.finish_weight_update() - self.model_runner.reset_lora_state() + if vllm_version_is("0.33.0"): + self.model_runner.reset_lora_state() + else: + self.model_runner.on_target_weights_replaced() self._weight_update_active = False def shutdown(self) -> None: @@ -457,6 +462,13 @@ def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks + def _get_gpu_memory_utilization(self) -> float: + utilization = self.cache_config.gpu_memory_utilization + # vLLM #56492 defers the default until the KV allocation mode is known. + if utilization is None: + return self.cache_config.resolved_gpu_memory_utilization + return utilization + def _init_device(self): # vLLM v0.24.0 (PR #45026) removed automatic per-process device # isolation for DP workers. Mirror gpu_worker.py::init_device: @@ -528,7 +540,7 @@ def _init_device(self): # take current memory snapshot self.init_snapshot = MemorySnapshot(device=device) - self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization + self.requested_memory = self.init_snapshot.total_memory * self._get_gpu_memory_utilization() if self.init_snapshot.free_memory < self.requested_memory: GiB = lambda b: round(b / GiB_bytes, 2) raise ValueError( @@ -536,7 +548,7 @@ def _init_device(self): f"({GiB(self.init_snapshot.free_memory)}/" f"{GiB(self.init_snapshot.total_memory)} GiB) on startup " f"is less than desired GPU memory utilization " - f"({self.cache_config.gpu_memory_utilization}, " + f"({self._get_gpu_memory_utilization()}, " f"{GiB(self.requested_memory)} GiB). Decrease GPU memory " f"utilization or reduce GPU memory used by other processes." ) @@ -1028,7 +1040,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes: f"({format_gib(self.init_snapshot.free_memory)}/" f"{format_gib(self.init_snapshot.total_memory)} GiB) on startup. " f"Desired GPU memory utilization is " - f"({self.cache_config.gpu_memory_utilization}, " + f"({self._get_gpu_memory_utilization()}, " f"{format_gib(self.requested_memory)} GiB). " f"Actual usage: {format_gib(self.model_runner.model_memory_usage)} GiB " f"for weights, {format_gib(self.peak_activation_memory)} GiB for peak " @@ -1439,7 +1451,11 @@ def get_encoder_cudagraph_stats(self) -> dict[str, Any] | None: @runtime_guard_idle_step def execute_dummy_batch(self) -> None: self.log_memory_stats() - num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1) + num_tokens = ( + getattr(self.model_runner, "uniform_decode_query_len", 1) + if vllm_version_is("0.33.0") + else self.model_runner.decode_query_len + ) # Idle DP worker RPCs also need the current config for replay metadata. with set_current_vllm_config(self.vllm_config): self.model_runner._dummy_run(num_tokens, uniform_decode=True, skip_gdn_state_update=True)