[Model][Feature] Support D3 decision models on Ascend

.github/workflows/scripts/estimated_times.yaml

@@ -48,6 +48,7 @@ estimated_times:
   tests/e2e/pull_request/one_card/model_runner_v2/test_uva.py: 80
   tests/e2e/pull_request/one_card/model_runner_v2/test_qwen3_8B.py: 350
   tests/e2e/pull_request/one_card/pooling/test_classification.py: 280
+  tests/e2e/pull_request/one_card/pooling/test_d3.py: 900
   tests/e2e/pull_request/one_card/pooling/test_embedding.py: 600
   tests/e2e/pull_request/one_card/pooling/test_scoring.py: 520
   tests/e2e/pull_request/one_card/spec_decode/test_dflash.py: 235

docs/source/user_guide/feature_guide/context_parallel.md

@@ -110,6 +110,7 @@ For either method, remove `--enforce-eager` and add the following option to use
   Eagle3 with GQA models, and DSpark with SFA, DeepSeek-V4 DSA, and DeepSeek-V4.1 DSA models.
 - PCP does not support `FULL` CUDA graphs; use eager mode, `PIECEWISE`, `FULL_AND_PIECEWISE`, or `FULL_DECODE_ONLY`.
 - Encoder-decoder models and LoRA are not supported with PCP.
+- Pooling models with bidirectional (`ENCODER_ONLY`) attention layers, such as Decision 3.0 (D3) checkpoints, are not supported with PCP. Pooling models default to the V1 model runner, which rejects PCP; if `VLLM_USE_V2_MODEL_RUNNER=1` forces the V2 runner, the attention backend fails at startup because it does not gather the other ranks' K/V.
 - PCP and DCP can be stacked for MLA, SFA, and GQA. See
   [Combining PCP and DCP](#combining-pcp-and-dcp) for the supported sizes.
 - Adaptive verification is not supported with PCP speculative decoding.

release_notes/fragments/19804.yaml

@@ -0,0 +1,4 @@
+pr: 19804
+entries:
+  - category: feature
+    text: Added an Ascend adapter for upstream Decision 3.0 checkpoints, using vLLM's native classification and readout sidecar paths. D3 requests use the upstream System One API and require a compatible vLLM revision with native D3 support.

release_notes/fragments/19919.yaml

@@ -0,0 +1,6 @@
+pr: 19919
+entries:
+  - category: bugfix
+    text: Fixed pooling models with bidirectional (ENCODER_ONLY) full-attention layers in hybrid architectures, such as Qwen3.5-based models, running causal attention on Ascend and returning incorrect results.
+  - category: bugfix
+    text: Pooling models with bidirectional (ENCODER_ONLY) attention layers now fail at startup when prefill context parallelism is enabled, instead of silently attending only to the rank-local tokens.

requirements-dev.txt

@@ -1,5 +1,6 @@
 -r requirements-lint.txt
 -r requirements.txt
+accelerate
 modelscope<1.38
 openai
 pytest >= 6.0,<9.0.0

tests/e2e/pull_request/one_card/pooling/test_d3.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib.util
+import json
+from typing import Any
+
+import pytest
+import requests
+import torch
+from transformers import AutoModel
+from vllm.utils.network_utils import get_open_port
+
+from tests.e2e.conftest import RemoteOpenAIServer, VllmRunner, cleanup_dist_env_and_memory
+
+# Probe for the feature instead of a version: native D3 (vLLM #61565, #61778) is in neither v0.31.0 nor v0.32.0.
+pytestmark = pytest.mark.skipif(
+    importlib.util.find_spec("vllm.model_executor.models.d3") is None,
+    reason="D3 needs a vLLM revision with native D3 support",
+)
+
+MODEL = "vllm-sr/d3-edge"
+# Allow for numerical differences between the Ascend and reference runtimes.
+SOFTMAX_ATOL = 0.03
+# These tests only send text, so media inputs are switched off. Keep every D3 run on one setting: the torch.compile
+# cache key of vLLM ignores `limit_mm_per_prompt`, and a graph compiled with the other setting crashes graph-mode
+# startup when it is reused.
+NO_MEDIA = {"image": 0, "video": 0}
+STATE = "A customer reports a damaged delivery and asks for a replacement."
+QUESTIONS: dict[str, dict[str, Any]] = {
+    "route": {
+        "type": "choice",
+        "instructions": "Which team should handle the request?",
+        "criteria": {
+            "returns": "Refunds and replacements",
+            "billing": "Payments and invoices",
+            "technical": "Product setup and faults",
+        },
+    },
+    "receipt": {
+        "type": "noul",
+        "instructions": "Does the customer ask for a replacement?",
+    },
+    "urgency": {
+        "type": "score",
+        "instructions": "How urgent is this request?",
+        "criteria": ["Routine", "Soon", "Today"],
+    },
+}
+
+
+def test_d3_pooling_matches_checkpoint_runtime():
+    reference_model = AutoModel.from_pretrained(MODEL, trust_remote_code=True)
+    runtime = reference_model.runtime
+    prepared = runtime.prepare(STATE, QUESTIONS)
+    runnable_keys = prepared.runnable
+    option_counts = [len(prepared.questions[key].keys) for key in runnable_keys]
+    prompt_ids = [prepared.sequences[key] for key in runnable_keys]
+    reference_logits = list(runtime.logits(prompt_ids, option_counts).cpu())
+    del reference_model, runtime
+    cleanup_dist_env_and_memory()
+
+    with VllmRunner(MODEL, runner="pooling", max_model_len=4096, limit_mm_per_prompt=NO_MEDIA) as runner:
+        outputs = runner.model.classify([{"prompt_token_ids": ids} for ids in prompt_ids], use_tqdm=False)
+
+    assert len(outputs) == len(option_counts)
+    for reference, output, option_count in zip(reference_logits, outputs, option_counts):
+        ascend_logits = torch.tensor(output.outputs.probs)
+        assert ascend_logits.shape == (255,)
+        assert torch.isfinite(ascend_logits).all()
+        torch.testing.assert_close(
+            torch.softmax(ascend_logits[:option_count].float(), dim=-1),
+            torch.softmax(reference[:option_count].float(), dim=-1),
+            atol=SOFTMAX_ATOL,
+            rtol=0,
+        )
+
+
+def test_d3_systemone_endpoint():
+    request = {
+        "model": MODEL,
+        "state": STATE,
+        "questions": QUESTIONS,
+    }
+    host = "127.0.0.1"
+    port = get_open_port()
+    server_args = [
+        "--host",
+        host,
+        "--port",
+        str(port),
+        "--max-model-len",
+        "4096",
+        "--limit-mm-per-prompt",
+        json.dumps(NO_MEDIA),
+    ]
+    with RemoteOpenAIServer(MODEL, server_args, server_host=host, server_port=port, auto_port=False) as server:
+        response = requests.post(
+            server.url_for("v1", "systemone"),
+            json=request,
+            timeout=600,
+        )
+
+    assert response.status_code == 200, response.text
+    payload = response.json()
+    assert payload["object"] == "structured_decision"
+    assert set(payload["answers"]) == set(QUESTIONS)
+
+    choice = payload["answers"]["route"]
+    assert choice["type"] == "choice"
+    assert choice["choice"] in QUESTIONS["route"]["criteria"]
+    assert set(choice["probabilities"]) == set(QUESTIONS["route"]["criteria"])
+
+    noul = payload["answers"]["receipt"]
+    assert noul["type"] == "noul"
+    assert set(noul["probabilities"]) == {"yes", "no"}
+    assert noul["noul"] == pytest.approx(noul["probabilities"]["yes"])
+
+    score = payload["answers"]["urgency"]
+    assert score["type"] == "score"
+    assert set(score["probabilities"]) == {"0", "1", "2"}
+
+    for answer in payload["answers"].values():
+        probabilities = list(answer["probabilities"].values())
+        assert all(torch.isfinite(torch.tensor(probabilities)))
+        assert all(probability >= 0 for probability in probabilities)
+        assert sum(probabilities) == pytest.approx(1.0, abs=1e-5)
+        assert 0 <= answer["confidence"] <= 1

tests/ut/attention/test_attention_v1.py

@@ -851,6 +851,75 @@ def test_forward_no_attn_metadata(self):
 
         assert output.shape == (10, 8 * 64)
 
+    def _build_impl(self, attn_type):
+        return AscendAttentionBackendImpl(
+            num_heads=8,
+            head_size=64,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=attn_type,
+            kv_sharing_target_layer_name=None,
+        )
+
+    def test_encoder_only_attention_rejects_prefill_context_parallel(self):
+        """The encoder attention path does not gather PCP K/V, so PCP must fail fast instead of being wrong."""
+        self.mock_vllm_config.parallel_config.prefill_context_parallel_size = 2
+
+        with self.assertRaisesRegex(NotImplementedError, "prefill context parallelism"):
+            self._build_impl(attn_module.AttentionType.ENCODER_ONLY)
+
+    def test_encoder_only_attention_without_pcp_is_accepted(self):
+        self.mock_vllm_config.parallel_config.prefill_context_parallel_size = 1
+
+        impl = self._build_impl(attn_module.AttentionType.ENCODER_ONLY)
+
+        assert impl.attn_type == attn_module.AttentionType.ENCODER_ONLY
+
+    def test_decoder_attention_with_pcp_is_not_rejected_by_the_encoder_only_guard(self):
+        self.mock_vllm_config.parallel_config.prefill_context_parallel_size = 2
+
+        impl = self._build_impl(attn_module.AttentionType.DECODER)
+
+        assert impl.attn_type == attn_module.AttentionType.DECODER
+
+    def _pooling_forward(self, attn_type):
+        """Run forward() for a pooling request whose metadata still says causal."""
+        query = torch.randn(10, 8, 64)
+        key = torch.randn(10, 8, 64)
+        value = torch.randn(10, 8, 64)
+        output = torch.zeros_like(query)
+        metadata = MagicMock()
+        metadata.model_runner_type = "pooling"
+        metadata.causal = True
+        metadata.uses_pcp_sharded_layout = False
+        self.impl.attn_type = getattr(attn_module.AttentionType, attn_type)
+        encoder_out = torch.ones(10, 8, 64)
+        with (
+            patch.object(self.impl, "reshape_and_cache", return_value=(query, key, value, output)),
+            patch.object(self.impl, "_forward_encoder_attention", return_value=encoder_out) as encoder,
+            patch.object(self.impl, "forward_impl", return_value=torch.zeros(10, 8, 64)) as causal_path,
+        ):
+            result = self.impl.forward(self.layer_no_quant, query, key, value, None, metadata, output)
+        return result, encoder, causal_path
+
+    def test_forward_pooling_encoder_only_is_bidirectional_despite_causal_metadata(self):
+        """ENCODER_ONLY layers (e.g. D3 full-attention layers in a hybrid model) must not use the causal path."""
+        result, encoder, causal_path = self._pooling_forward("ENCODER_ONLY")
+
+        encoder.assert_called_once()
+        causal_path.assert_not_called()
+        assert torch.equal(result, torch.ones(10, 8, 64))
+
+    def test_forward_pooling_decoder_keeps_causal_path(self):
+        result, encoder, causal_path = self._pooling_forward("DECODER")
+
+        encoder.assert_not_called()
+        causal_path.assert_called_once()
+
     @patch("torch_npu.npu_scatter_pa_kv_cache")
     @patch("torch_npu.npu_fused_infer_attention_score")
     @patch("vllm_ascend.ascend_forward_context.get_forward_context")

tests/ut/models/test_d3_registration.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project
+
+from unittest.mock import patch
+
+from vllm_ascend import models
+from vllm_ascend.models.d3 import AscendD3ForSequenceClassification
+
+
+def test_d3_sequence_classifier_uses_ascend_adapter():
+    registrations = []
+    with patch.object(
+        models.ModelRegistry,
+        "register_model",
+        side_effect=lambda name, path: registrations.append((name, path)),
+    ):
+        models.register_model()
+
+    assert (
+        "D3ForSequenceClassification",
+        "vllm_ascend.models.d3:AscendD3ForSequenceClassification",
+    ) in registrations
+
+
+def test_d3_skips_readout_weight_in_backbone_pass():
+    assert AscendD3ForSequenceClassification.is_unused_checkpoint_weight("weight")
+    assert not AscendD3ForSequenceClassification.is_unused_checkpoint_weight("language_model.model.embed_tokens.weight")

tests/ut/test_platform.py

@@ -280,6 +280,16 @@ def test_sfa_dcp_validation_only_bypasses_separate_draft(model_role):
             enable_sfa.assert_called_once_with(config)
 
 
+def test_pcp_is_rejected_on_the_v1_model_runner():
+    """Pooling models default to the V1 runner (see mrv2_utils), where PCP is not supported."""
+    config = SimpleNamespace(
+        use_v2_model_runner=False, parallel_config=SimpleNamespace(prefill_context_parallel_size=2)
+    )
+
+    with pytest.raises(ValueError, match="PCP .* is not supported by vLLM Ascend"):
+        _validate_parallel_config(config)
+
+
 @pytest.mark.parametrize("device_type", [AscendDeviceType.A2, AscendDeviceType.A3, AscendDeviceType.A5])
 @pytest.mark.parametrize("enable_sfa_c8", [False, True])
 def test_sfa_dcp_c8_hardware_validation(device_type, enable_sfa_c8):

vllm_ascend/attention/attention_v1.py

@@ -579,6 +579,15 @@ def __init__(
         **kwargs,
     ) -> None:
         self.vllm_config = get_current_vllm_config()
+        if (
+            attn_type == AttentionType.ENCODER_ONLY
+            and self.vllm_config.parallel_config.prefill_context_parallel_size > 1
+        ):
+            raise NotImplementedError(
+                "Bidirectional (ENCODER_ONLY) attention does not support prefill context parallelism (PCP): each "
+                "PCP rank only holds its slice of the prompt and the encoder attention path does not gather the "
+                "other ranks' K/V, so the result would silently be wrong. Set --prefill-context-parallel-size to 1."
+            )
         self.use_pcp = (
             type(self) is AscendAttentionBackendImpl
             and self.vllm_config.parallel_config.prefill_context_parallel_size > 1
@@ -1405,8 +1414,11 @@ def forward(
                 query, key, value, output_padded = self.reshape_and_cache(
                     query, key, value, kv_cache, attn_metadata, output
                 )
-        # pooling model branch
-        if attn_metadata.model_runner_type == "pooling" and not attn_metadata.causal:
+        # pooling model branch. ENCODER_ONLY layers are bidirectional even when the metadata of a
+        # hybrid model still says causal (they never go through the EncoderOnlyAttention wrapper).
+        if attn_metadata.model_runner_type == "pooling" and (
+            not attn_metadata.causal or self.attn_type == AttentionType.ENCODER_ONLY
+        ):
             attn_output = self._forward_encoder_attention(query, key, value, attn_metadata, output)
             output[:num_tokens] = attn_output[:num_tokens]
             return output

vllm_ascend/models/__init__.py

@@ -95,6 +95,10 @@ def register_model():
         "Qwen3_5MoeForConditionalGeneration",
         "vllm_ascend.models.qwen3_5_vl:AscendQwen3_5MoeForConditionalGeneration",
     )
+    ModelRegistry.register_model(
+        "D3ForSequenceClassification",
+        "vllm_ascend.models.d3:AscendD3ForSequenceClassification",
+    )
     ModelRegistry.register_model(
         "Eagle3LlamaForCausalLM", "vllm_ascend.models.llama_eagle3:AscendEagle3LlamaForCausalLM"
     )