diff --git a/tests/ut/core/test_recompute_scheduler.py b/tests/ut/core/test_recompute_scheduler.py index 70c8c703c65b..f3bd5ce7a132 100644 --- a/tests/ut/core/test_recompute_scheduler.py +++ b/tests/ut/core/test_recompute_scheduler.py @@ -21,6 +21,17 @@ ) +@pytest.fixture(autouse=True) +def _oproj_tp_config(monkeypatch): + # _preempt_or_recompute reads the live AscendConfig singleton; default it to off. + finegrained_tp_config = SimpleNamespace(oproj_tensor_parallel_size=0) + monkeypatch.setattr( + "vllm_ascend.core.recompute_scheduler.get_ascend_config", + lambda: SimpleNamespace(finegrained_tp_config=finegrained_tp_config), + ) + return finegrained_tp_config + + def _make_preempt_scheduler(*, connector=None): scheduler = RecomputeScheduler.__new__(RecomputeScheduler) scheduler.connector = connector @@ -140,6 +151,38 @@ def test_preempt_offload_failure_sends_request_back_to_p(): scheduler.finish_requests.assert_called_once() +def test_preempt_offload_failure_aborts_under_oproj_tp(_oproj_tp_config): + _oproj_tp_config.oproj_tensor_parallel_size = 2 + connector = MagicMock() + connector.update_state_before_preempt.return_value = False + scheduler = _make_preempt_scheduler(connector=connector) + request = SimpleNamespace( + request_id="req-1", + client_index=0, + num_computed_tokens=17, + ) + scheduler.finish_requests = MagicMock(return_value=[request]) + + with ( + patch.object(Scheduler, "_preempt_request") as upstream_preempt, + patch("vllm_ascend.core.recompute_scheduler.logger.error") as error, + ): + locally_preempted = scheduler._preempt_or_recompute( + request, + 1.5, + ) + + assert not locally_preempted + error.assert_called_once() + upstream_preempt.assert_not_called() + scheduler.finish_requests.assert_called_once_with( + "req-1", + RequestStatus.FINISHED_ABORTED, + ) + # Not recorded for recomputation: the request must not be sent back to P. + assert scheduler._recomputed_reqs == [] + + def test_reset_preemption_skips_offload(): connector = MagicMock() scheduler = _make_preempt_scheduler(connector=connector) diff --git a/tests/ut/test_ascend_config.py b/tests/ut/test_ascend_config.py index ab29e5fd7c92..c6b89dd574fd 100644 --- a/tests/ut/test_ascend_config.py +++ b/tests/ut/test_ascend_config.py @@ -24,6 +24,7 @@ from vllm.config import KVTransferConfig from vllm.config import VllmConfig as _VllmConfig +from vllm.config.compilation import CUDAGraphMode from tests.ut.base import TestBase from tests.ut.kvpp_utils import make_kvpp_config @@ -967,6 +968,76 @@ def test_finegrained_tp_config_rejects_negative_size(self): with self.assertRaisesRegex(ValueError, "lmhead_tensor_parallel_size must be non-negative"): FinegrainedTPConfig(lmhead_tensor_parallel_size=-1) + def _oproj_tp_vllm_config( + self, + max_num_batched_tokens=8192, + max_num_seqs=256, + num_speculative_tokens=0, + max_cudagraph_capture_size=512, + cudagraph_capture_sizes=None, + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + prefill_context_parallel_size=1, + ): + speculative_config = None + if num_speculative_tokens: + speculative_config = SimpleNamespace(num_speculative_tokens=num_speculative_tokens) + return SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=1, + data_parallel_size=8, + prefill_context_parallel_size=prefill_context_parallel_size, + ), + compilation_config=SimpleNamespace( + cudagraph_mode=cudagraph_mode, + max_cudagraph_capture_size=max_cudagraph_capture_size, + cudagraph_capture_sizes=cudagraph_capture_sizes, + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=max_num_batched_tokens, max_num_seqs=max_num_seqs), + speculative_config=speculative_config, + kv_transfer_config=SimpleNamespace(is_kv_consumer=True), + model_config=SimpleNamespace(is_moe=True), + ) + + def test_oproj_tp_requires_graph_mode(self): + config = FinegrainedTPConfig(oproj_tensor_parallel_size=2) + # VllmConfig.__post_init__ normalizes enforce_eager into NONE, so this + # single check covers both spellings of "no graph mode". + with self.assertRaisesRegex(AssertionError, "only supported in graph mode"): + config._validate_preconditions(self._oproj_tp_vllm_config(cudagraph_mode=CUDAGraphMode.NONE)) + + def test_oproj_tp_rejects_pcp(self): + config = FinegrainedTPConfig(oproj_tensor_parallel_size=2) + with self.assertRaisesRegex(AssertionError, "not supported with prefill_context_parallel_size"): + config._validate_preconditions(self._oproj_tp_vllm_config(prefill_context_parallel_size=2)) + + def test_oproj_tp_size_one_skips_the_checks(self): + # Size 1 requests no split: no exchange groups to align, so the preconditions do not apply. + config = FinegrainedTPConfig(oproj_tensor_parallel_size=1) + config._validate_preconditions(self._oproj_tp_vllm_config(cudagraph_mode=CUDAGraphMode.NONE)) + self.assertEqual(config.oproj_tensor_parallel_size, 1) + + def test_oproj_tp_capture_bound_check(self): + config = FinegrainedTPConfig(oproj_tensor_parallel_size=2) + config._validate_preconditions(self._oproj_tp_vllm_config()) + # max_num_batched_tokens can cap the step below the capture bound. + config._validate_preconditions(self._oproj_tp_vllm_config(max_num_batched_tokens=512)) + self.assertEqual(config.oproj_tensor_parallel_size, 2) + # 300 reqs x decode_query_len 2 (spec window) = 600 > 512: disabled with a warning. + config._validate_preconditions(self._oproj_tp_vllm_config(max_num_seqs=300, num_speculative_tokens=1)) + self.assertEqual(config.oproj_tensor_parallel_size, 0) + # An explicit capture size that covers the step keeps the knob on. + config = FinegrainedTPConfig(oproj_tensor_parallel_size=2) + config._validate_preconditions( + self._oproj_tp_vllm_config(max_num_seqs=300, num_speculative_tokens=1, max_cudagraph_capture_size=1024) + ) + self.assertEqual(config.oproj_tensor_parallel_size, 2) + # Before _set_cudagraph_sizes backfills it, an explicit sizes list is the bound. + config = FinegrainedTPConfig(oproj_tensor_parallel_size=2) + config._validate_preconditions( + self._oproj_tp_vllm_config(max_cudagraph_capture_size=None, cudagraph_capture_sizes=[8, 16, 512]) + ) + self.assertEqual(config.oproj_tensor_parallel_size, 2) + def test_eplb_config_int_field_lax(self): cfg = EplbConfig(eplb_policy_type="2") self.assertEqual(cfg.eplb_policy_type, 2) diff --git a/tests/ut/test_platform.py b/tests/ut/test_platform.py index 7c8bf046759b..49f3746db29b 100644 --- a/tests/ut/test_platform.py +++ b/tests/ut/test_platform.py @@ -165,6 +165,7 @@ def mock_vllm_ascend_config(): mock_ascend_config.ascend_compilation_config.enable_npugraph_ex = False mock_ascend_config.ascend_fusion_config = None mock_ascend_config.scheduler_config.recompute_scheduler_enable = False + mock_ascend_config.finegrained_tp_config.oproj_tensor_parallel_size = 0 mock_ascend_config.scheduler_config.enable_balance_scheduling = False mock_ascend_config.scheduler_config.batch_job_sched_config.enabled = False mock_ascend_config.mc2_comm_alg = "" @@ -583,6 +584,21 @@ def test_check_and_update_config_validates_indexer_before_worker_start( mock_validate_indexer.assert_called_once_with(vllm_config) + def test_check_ascend_config_oproj_tp_requires_offload_connector(self): + from vllm_ascend.platform import _check_ascend_config + + vllm_config = TestNPUPlatform.mock_vllm_config() + ascend_config = TestNPUPlatform.mock_vllm_ascend_config() + ascend_config.finegrained_tp_config.oproj_tensor_parallel_size = 2 + + # The base mock carries no kv_transfer_config: a real split must fail closed. + with pytest.raises(AssertionError, match="PreemptOffloadConnector"): + _check_ascend_config(vllm_config, ascend_config) + + # Size 1 exchanges nothing across ranks and stays exempt. + ascend_config.finegrained_tp_config.oproj_tensor_parallel_size = 1 + _check_ascend_config(vllm_config, ascend_config) + def test_apply_config_platform_defaults_skips_when_scheduler_max_num_seqs_is_missing(self): vllm_config = TestNPUPlatform.mock_vllm_config() vllm_config.compilation_config.max_cudagraph_capture_size = None diff --git a/tests/ut/worker/test_model_runner_v2.py b/tests/ut/worker/test_model_runner_v2.py index 5851bafe3915..cbbdcc9f0c75 100644 --- a/tests/ut/worker/test_model_runner_v2.py +++ b/tests/ut/worker/test_model_runner_v2.py @@ -33,6 +33,8 @@ def _make_runner(need_timing: bool = True): runner.attn_groups = [] runner.adaptive_verification = None runner.use_fia = False + # Set by NPUModelRunner.__init__ on real instances. + runner._oproj_tp_requires_graph = False return runner @@ -465,6 +467,14 @@ def _parent_init(self, vllm_config, device, *, full_graph=False, speculative=Fal def test_init_without_spec_pp(): vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(enable_eplb=False)) ascend_config = SimpleNamespace(eplb_config=SimpleNamespace(load_collection_phase="all")) + + # Complete the fake with the fields NPUModelRunner reads (mirrors FinegrainedTPConfig). + ascend_config.finegrained_tp_config = SimpleNamespace( + oproj_tensor_parallel_size=0, + lmhead_tensor_parallel_size=0, + embedding_tensor_parallel_size=0, + mlp_tensor_parallel_size=0, + ) with ( patch("vllm_ascend.worker.v2.model_runner.get_ascend_config", return_value=ascend_config), patch("vllm_ascend.worker.v2.model_runner.set_potential_max_tokens"), @@ -501,6 +511,14 @@ def test_init_without_spec_pp(): def test_init_spec_pp_full_graph_and_speculator(): vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(enable_eplb=True)) ascend_config = SimpleNamespace(eplb_config=SimpleNamespace(load_collection_phase="decode")) + + # Complete the fake with the fields NPUModelRunner reads (mirrors FinegrainedTPConfig). + ascend_config.finegrained_tp_config = SimpleNamespace( + oproj_tensor_parallel_size=0, + lmhead_tensor_parallel_size=0, + embedding_tensor_parallel_size=0, + mlp_tensor_parallel_size=0, + ) spec_pp = SimpleNamespace(needs_aux_hidden_states=True) speculator = SimpleNamespace() with ( diff --git a/tests/ut/worker/test_model_runner_v2_finegrained_tp.py b/tests/ut/worker/test_model_runner_v2_finegrained_tp.py index b4cbe29e8da3..1dba5a91a50a 100644 --- a/tests/ut/worker/test_model_runner_v2_finegrained_tp.py +++ b/tests/ut/worker/test_model_runner_v2_finegrained_tp.py @@ -1,9 +1,11 @@ -"""Unit tests for lmhead TP support in the Ascend V2 model runner. +"""Unit tests for fine-grained TP support in the Ascend V2 model runner. Pure-mock tests (CPU tensors, no NPU): they lock the runner-side pad/trim -contract of sample()/_dummy_run and guard the copied dispatch tail with a -canary that compares it call-by-call against upstream GPUModelRunner.sample. -Collective behavior of the LM head itself is validated on real hardware. +contract of sample()/_dummy_run (lmhead TP), guard the copied dispatch tail +with a canary that compares it call-by-call against upstream +GPUModelRunner.sample, and pin the o_proj TP graph-mode guard that turns an +eagerly dispatched step into an explicit error. Collective behavior itself +is validated on real hardware. """ from types import SimpleNamespace @@ -11,6 +13,7 @@ import pytest import torch +from vllm.config.compilation import CUDAGraphMode from vllm.v1.worker.gpu.eplb_utils import step_eplb_after from vllm.v1.worker.gpu.model_runner import GPUModelRunner from vllm.v1.worker.gpu.sample.sampler import Sampler @@ -238,3 +241,14 @@ def compute_logits(inputs): torch.testing.assert_close(dummy_input, hidden_states[torch.zeros(16, dtype=torch.long)]) else: runner.model.compute_logits.assert_not_called() + + +def test_oproj_tp_guard_contract(): + runner = object.__new__(NPUModelRunner) + runner._oproj_tp_requires_graph = False + NPUModelRunner._check_oproj_tp_graph_step(runner, CUDAGraphMode.NONE) + runner._oproj_tp_requires_graph = True + with pytest.raises(RuntimeError, match="captured graph"): + NPUModelRunner._check_oproj_tp_graph_step(runner, CUDAGraphMode.NONE) + NPUModelRunner._check_oproj_tp_graph_step(runner, CUDAGraphMode.FULL_DECODE_ONLY) + NPUModelRunner._check_oproj_tp_graph_step(runner, CUDAGraphMode.FULL) diff --git a/vllm_ascend/ascend_config.py b/vllm_ascend/ascend_config.py index 5073bd56dceb..f4cff402197c 100644 --- a/vllm_ascend/ascend_config.py +++ b/vllm_ascend/ascend_config.py @@ -689,16 +689,14 @@ def derive_and_validate(self, vllm_config: VllmConfig) -> AscendConfig: str(vc.scheduler_config.max_num_batched_tokens), ) - # finegrained_tp requires recompute_scheduler - if ( + finegrained_tp_enabled = ( self.finegrained_tp_config.oproj_tensor_parallel_size > 0 or self.finegrained_tp_config.embedding_tensor_parallel_size > 0 - ) and not self.scheduler_config.recompute_scheduler_enable: + ) + if finegrained_tp_enabled and not self.scheduler_config.recompute_scheduler_enable: raise AssertionError( - "oproj_tensor_parallel_size / embedding_tensor_parallel_size " - "require recompute_scheduler_enable=true: their cross-DP HCCL " - "collectives need uniform num_tokens across DP ranks, which is " - "only guaranteed when the recompute scheduler is enabled." + "oproj_tensor_parallel_size / embedding_tensor_parallel_size require " + "recompute_scheduler_enable=true: it keeps decode-node steps decode-shaped.", ) # enable_fused_mc2 enum + MiniMax mutex + multistream auto-disable @@ -1177,32 +1175,59 @@ def _validate_sizes(self): return self def _validate_preconditions(self, vllm_config: Any): + # Local import to avoid a circular import during platform resolution. + from vllm.config.compilation import CUDAGraphMode + vc = vllm_config enabled_configs = [] - if self.oproj_tensor_parallel_size > 0: - enabled_configs.append(f"oproj_tensor_parallel_size={self.oproj_tensor_parallel_size}") - # wo_a/wo_b are sharded solely by the OTP group (which splits DP, - # orthogonal to the standard TP group), but _forward_o_proj reshapes - # the attention output with n_local_groups = n_groups // tp_size - # (standard TP). When tp_size > 1 the weight-shard and input-shard - # operate on different axes of the rank grid and no longer align, - # so oproj TP currently requires standard tp_size == 1. + if self.oproj_tensor_parallel_size > 1: + # _forward_o_proj reshapes with n_local_groups = n_groups // tp_size (standard TP), + # which misaligns with the OTP weight shard (DP axis) when tp_size > 1. if vc.parallel_config.tensor_parallel_size > 1: raise AssertionError( "oproj_tensor_parallel_size currently requires " "tensor_parallel_size == 1, got " f"{vc.parallel_config.tensor_parallel_size}." ) - # The static all_to_all / reduce_scatter exchange buffers used by - # _forward_o_proj are sized for graph replay and require ACL graph - # capture; dummy_run does not run the entire attention module in - # eager mode, so o_proj tp split can only be used in graph mode. - if vc.model_config and vc.model_config.enforce_eager: + # Graph dispatch is the only lane that aligns DP token counts (eager keeps per-rank counts). + if vc.compilation_config.cudagraph_mode == CUDAGraphMode.NONE: raise AssertionError("oproj_tensor_parallel_size is only supported in graph mode") if vc.kv_transfer_config is None or not vc.kv_transfer_config.is_kv_consumer: raise AssertionError( "oproj_tensor_parallel_size is only supported in pd scenario and can only be used in D node." ) + # PCP's dispatch recomputes num_tokens per rank, breaking the group-uniform step size. + if vc.parallel_config.prefill_context_parallel_size > 1: + raise AssertionError( + "oproj_tensor_parallel_size is not supported with prefill_context_parallel_size > 1." + ) + # decode_query_len mirrors _get_default_max_cudagraph_capture_size in platform.py. + decode_query_len = 1 + speculative_config = vc.speculative_config + if speculative_config and speculative_config.num_speculative_tokens: + decode_query_len += speculative_config.num_speculative_tokens + max_step = min( + vc.scheduler_config.max_num_batched_tokens, vc.scheduler_config.max_num_seqs * decode_query_len + ) + capture_bound = vc.compilation_config.max_cudagraph_capture_size + # An explicit sizes list is the bound until _set_cudagraph_sizes backfills the capture max. + if capture_bound is None: + capture_sizes = vc.compilation_config.cudagraph_capture_sizes + capture_bound = max(capture_sizes) if capture_sizes else None + # A step beyond the capture bound dispatches to eager and desyncs the cross-DP collectives. + if capture_bound is None or capture_bound < max_step: + logger.warning( + "Disabling oproj_tensor_parallel_size=%d: the largest cudagraph capture " + "size (%s) does not cover the largest possible step (%d tokens); an " + "oversized step would dispatch to eager and hang the cross-DP HCCL " + "collectives. Raise max_cudagraph_capture_size to re-enable it.", + self.oproj_tensor_parallel_size, + str(capture_bound), + max_step, + ) + self.oproj_tensor_parallel_size = 0 + else: + enabled_configs.append(f"oproj_tensor_parallel_size={self.oproj_tensor_parallel_size}") if self.lmhead_tensor_parallel_size > 0: enabled_configs.append(f"lmhead_tensor_parallel_size={self.lmhead_tensor_parallel_size}") if self.embedding_tensor_parallel_size > 0: diff --git a/vllm_ascend/core/recompute_scheduler.py b/vllm_ascend/core/recompute_scheduler.py index 27d36f41c9f7..c755b0ee4a60 100644 --- a/vllm_ascend/core/recompute_scheduler.py +++ b/vllm_ascend/core/recompute_scheduler.py @@ -46,7 +46,7 @@ DyntraLBPolicyMixin, print_scheduler_summary, ) -from vllm_ascend.utils import vllm_version_is +from vllm_ascend.utils import get_ascend_config, vllm_version_is @dataclass @@ -87,8 +87,8 @@ class RecomputeScheduler(Scheduler): This keeps a local copy of vLLM's schedule() only to pad the first decode request for stable Ascend speculative-decode graph shapes. Preempted KV is - offloaded when possible; otherwise the request is sent back to P to redo - prefill. + offloaded when possible; otherwise the request returns to P to redo + prefill, or is aborted under o_proj TP. """ prefill_capacity_bound: bool @@ -132,6 +132,17 @@ def _preempt_or_recompute( ) if not offloaded: + # Mirror the config gate: only a real split (size > 1) forbids the return to P. + if get_ascend_config().finegrained_tp_config.oproj_tensor_parallel_size > 1: + logger.error( + "KV offload failed with o_proj TP enabled; aborting the request instead " + "of returning it to P for recomputation: request_id=%s, " + "num_computed_tokens=%d", + request.request_id, + request.num_computed_tokens, + ) + self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + return False if not offload_raised: logger.warning( "KV offload was unavailable or failed before decode-side " diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index 5d06b57c6603..a0a9ff6f5a8d 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -1043,6 +1043,18 @@ def _check_ascend_config(vllm_config: VllmConfig, ascend_config) -> None: ) vllm_config.scheduler_config = recompute_scheduler_config + # Checked here, not in AscendConfig: MultiConnector children re-validate the config + # with per-child copies that cannot see sibling connectors. + kv_transfer_config = vllm_config.kv_transfer_config + offload_missing = kv_transfer_config is None or not kv_transfer_config.has_connector("PreemptOffloadConnector") + # Only a real split (size > 1) needs the offload guarantee, mirroring the runner gate. + if offload_missing and ascend_config.finegrained_tp_config.oproj_tensor_parallel_size > 1: + raise AssertionError( + "oproj_tensor_parallel_size requires PreemptOffloadConnector (via MultiConnector): " + "a preempted request must not return to the prefill node, whose recomputed KV " + "loses precision." + ) + def _validate_kv_load_failure_policy(vllm_config: VllmConfig) -> None: kv_transfer_config = vllm_config.kv_transfer_config diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 6681c120b8f2..d0ffd66290ce 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -127,6 +127,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): if spec_pp_support is not None and spec_pp_support.needs_aux_hidden_states: self.use_aux_hidden_state_outputs = True + # Only a real split (size > 1) exchanges across ranks, and graph dispatch keeps it aligned. + ftpc = self.ascend_config.finegrained_tp_config + self._oproj_tp_requires_graph = ftpc.oproj_tensor_parallel_size > 1 and self.dp_size > 1 + self.use_aclgraph = ( self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE and ( @@ -389,6 +393,16 @@ def gather_batch_req_state(self, scheduler_output: SchedulerOutput, dummy_run: b _PCP_DISPATCH_NUM_TOKENS.set(num_tokens) return batch_state, uniform_token_count + def _check_oproj_tp_graph_step(self, cg_mode: CUDAGraphMode) -> None: + # Eager dispatch keeps per-rank token counts; the cross-DP o_proj exchange would desync. + if self._oproj_tp_requires_graph and cg_mode == CUDAGraphMode.NONE: + raise RuntimeError( + "o_proj TP requires every step on a captured graph: this step dispatched " + "to eager, which desyncs the cross-DP HCCL collectives (mixed or oversized " + "batch, a request-arrival step misclassified as prefill, or a full prefill " + "scheduled locally — a request sent directly to the decode node)." + ) + def prepare_inputs( # type: ignore[misc] self, scheduler_output: SchedulerOutput, @@ -399,6 +413,7 @@ def prepare_inputs( # type: ignore[misc] npu attention backends need seq_lens_cpu to work. so we need to prepare seq_lens_cpu here. """ + self._check_oproj_tp_graph_step(batch_desc.cg_mode) num_tokens = batch_req_state.num_tokens num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens) assert num_tokens > 0