[Feature][MRV2] Support o_proj TP in model runner v2 (graph-mode scope)

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)

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)

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

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 (

tests/ut/worker/test_model_runner_v2_finegrained_tp.py

@@ -1,16 +1,19 @@
-"""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
 from unittest.mock import MagicMock, create_autospec, patch
 
 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)

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:

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 "