[BugFix](pcp): align decode classification with vLLM reorder logic (vllm-project#10580)

## What this PR does / why we need it?

With PCP enabled and speculative decoding (e.g. MTP),
`PCPManager.init_batch_info()` previously marked a request as decode
whenever `num_scheduled_tokens <= decode_threshold`. That is
insufficient for chunked prefill: a request can still be prefilling
while scheduling only a small token chunk (≤ `mtp + 1`), and it would be
misclassified as decode.

Misclassified requests were routed through the PCP **decode** attention
path with inconsistent metadata (e.g. `actual_seq_lengths_q` vs local
`Q_S`), which triggered `aclnnFusedInferAttentionScoreV3` failures
(error `561002`).

This PR aligns PCP decode/prefill classification with upstream vLLM
`reorder_batch_to_split_decodes_and_prefills`:

- Add `classify_decode_request_mask()` in `attention/utils.py`. A
request is treated as decode only if it has context
(`num_computed_tokens > 0`), is below the decode threshold, and has
finished prefilling (`num_computed_tokens >= num_prompt_tokens`).
- Update `PCPManager.init_batch_info()` to use this mask instead of a
scheduled-token threshold alone, and pass `num_computed_tokens` /
`num_prompt_tokens` from model runners.
- Extend `split_decodes_and_prefills()`:
- Default `treat_short_extends_as_decodes=False`, so short extends are
prefills by default.
  - Honor `is_prefilling` when short extends are not treated as decodes.
- Add a `num_reqs == 0` guard and PyTorch tensor input support in
`classify_decode_request_mask()`.
- Set `is_prefilling` in spec-decode graph-capture dummy-run paths
(`llm_base_proposer`, `dflash_proposer`) to avoid startup assertion
failures.
- Correct stale E2E golden outputs in `test_accuracy.py` for DSV2 PCP
and Qwen3 PCP/DCP cases after re-verifying greedy outputs on NPU.

## Does this PR introduce any user-facing change?

No API or configuration changes. This is a bug fix for PCP + speculative
decoding workloads, especially with chunked prefill.

Users who previously hit decode/prefill misclassification and NPU
attention kernel errors should see correct routing and stable inference.
Behavior for true decode batches is unchanged.

## How was this patch tested?

**Unit tests**

- Added `test_split_decodes_short_extend_with_default_false` in
`tests/ut/worker/test_pcp_manager.py` to verify short extends with
`is_prefilling=True` are classified as prefills when
`decode_threshold=4` (MTP scenario).
- Updated existing PCP manager tests to pass `num_computed_tokens` and
`num_prompt_tokens` into `init_batch_info`.

**Manual reproduction**

- Reproduced the original failure with PCP + MTP + chunked prefill
(short prefill chunk misrouted to decode kernel).
- Verified the fix routes those chunks through the prefill path and
inference completes without `aclnnFusedInferAttentionScoreV3` error
`561002`.

**E2E**

- Re-ran four-card CP accuracy tests on NPU and updated
`DSV2_PCP_GOLDEN` / `QWEN3_GOLDEN` to match verified greedy outputs.

DeepSeek-V3.1-Terminus-w8a8-mtp-QuaRot test result

| dataset | version | metric | mode | vllm-api-stream-chat |
|----- | ----- | ----- | ----- | -----|
| gsm8k | d486ce | accuracy | gen | 94.20 |

Qwen3-30B-A3B-W8A8  test result
| dataset | version | metric | mode | vllm-api-stream-chat |
|----- | ----- | ----- | ----- | -----|
| gsm8k | d486ce | accuracy | gen | 92.80 |



**CI**

- Existing unit tests and lint checks.

- vLLM version: v0.22.1
- vLLM main:
vllm-project/vllm@967c5c3

---------

Signed-off-by: recky-c <ruiqicheng510@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: recky-c

tests/e2e/pull_request/four_card/context_parallel/test_accuracy.py

@@ -56,7 +56,7 @@
 
 DSV2_PCP_GOLDEN = [
     "The president of the United States is a man who is not only a liar, but",
-    "The capital of France is Paris.\nThe capital of the United States is",
+    "The capital of France is Paris.\nThe currency of France is the Euro",
 ]
 
 DSV2_DCP_GOLDEN = [
@@ -65,7 +65,7 @@
 ]
 
 QWEN3_GOLDEN = [
-    "The capital of France is Paris. Which of the",
+    "The capital of France is Paris. The capital of",
     "Hello, my name is Tom, I am 12 years old",
     "The president of United States is the head of state and",
 ]

tests/ut/ops/test_gdn_attn_builder.py

@@ -82,6 +82,16 @@ def create_common_attn_metadata(
     max_seq_len = int(seq_lens_cpu.max())
     context_lens = [batch_spec.seq_lens[i] - batch_spec.query_lens[i] for i in range(batch_spec.batch_size)]
     num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
+    # Mirror model_runner: is_prefilling = num_computed < num_prompt_tokens.
+    # Chunked prefills still have prompt tokens beyond num_computed; decodes do not.
+    num_prompt_tokens_cpu = torch.tensor(
+        [
+            context_lens[i] + batch_spec.query_lens[i] if batch_spec.query_lens[i] > 1 else context_lens[i]
+            for i in range(batch_spec.batch_size)
+        ],
+        dtype=torch.int32,
+    )
+    is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu
     max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size
     block_table_tensor = torch.arange(
         batch_spec.batch_size * max_blocks,
@@ -103,6 +113,7 @@ def create_common_attn_metadata(
         block_table_tensor=block_table_tensor,
         slot_mapping=slot_mapping,
         causal=True,
+        is_prefilling=is_prefilling,
     )
 
 

tests/ut/worker/test_pcp_manager.py

@@ -18,6 +18,7 @@
 import pytest
 import torch
 
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata, split_decodes_and_prefills
 from vllm_ascend.worker.pcp_utils import PCPManager
 
 
@@ -150,7 +151,14 @@ def test_generate_pcp_metadata_mla_tail_projection_indices(pcp_size, pcp_rank, q
 
     num_reqs = len(query_lens)
     num_scheduled_tokens = np.array(query_lens, dtype=np.int32)
-    pcp_manager.init_batch_info(num_scheduled_tokens, num_reqs)
+    num_computed_tokens = np.zeros(num_reqs, dtype=np.int32)
+    num_prompt_tokens = np.array(query_lens, dtype=np.int32)
+    pcp_manager.init_batch_info(
+        num_scheduled_tokens,
+        num_reqs,
+        num_computed_tokens,
+        num_prompt_tokens,
+    )
 
     input_batch = MagicMock()
     input_batch.num_reqs = num_reqs
@@ -251,7 +259,12 @@ def test_update_tokens_for_pcp_basic(
     input_batch.num_prompt_tokens = np.array(num_prompt_tokens, dtype=np.int32)
     arange_np = np.arange(10000)
     num_scheduled_tokens = np.array(tokens)
-    pcp_manager.init_batch_info(num_scheduled_tokens, num_reqs)
+    pcp_manager.init_batch_info(
+        num_scheduled_tokens,
+        num_reqs,
+        input_batch.num_computed_tokens_cpu,
+        input_batch.num_prompt_tokens,
+    )
     pcp_tokens_result, positions_result = pcp_manager.update_tokens_for_pcp(num_scheduled_tokens, arange_np)
 
     assert np.array_equal(pcp_tokens_result, expected_pcp_tokens), (
@@ -264,6 +277,39 @@ def test_update_tokens_for_pcp_basic(
     )
 
 
+def test_split_decodes_short_extend_with_default_false():
+    """Short extends should be treated as prefills by default."""
+    long_seq_metadata = MagicMock()
+    long_seq_metadata.query_lens_pcp_full_cpu = torch.tensor([3], dtype=torch.int32)
+    long_seq_metadata.max_query_len_pcp_full = 3
+
+    query_start_loc_cpu = torch.tensor([0, 2], dtype=torch.int32)
+    common_attn_metadata = AscendCommonAttentionMetadata(
+        query_start_loc=query_start_loc_cpu,
+        query_start_loc_cpu=query_start_loc_cpu,
+        seq_lens=torch.tensor([173], dtype=torch.int32),
+        num_reqs=1,
+        num_actual_tokens=2,
+        max_query_len=2,
+        max_seq_len=173,
+        block_table_tensor=torch.zeros((1, 1), dtype=torch.int32),
+        slot_mapping=torch.arange(2, dtype=torch.int32),
+        is_prefilling=torch.tensor([True]),
+        prefill_context_parallel_metadata=long_seq_metadata,
+    )
+
+    num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = split_decodes_and_prefills(
+        common_attn_metadata,
+        decode_threshold=4,
+        treat_short_extends_as_decodes=False,
+    )
+
+    assert num_decodes == 0
+    assert num_prefills == 1
+    assert num_decode_tokens == 0
+    assert num_prefill_tokens == 2
+
+
 # yapf: disable
 @pytest.mark.parametrize(
     "seq_lens, pcp_world_size, dcp_world_size, cp_kv_cache_interleave_size, target",
@@ -417,7 +463,20 @@ def test_generate_pcp_mtp_input(
     for i, token_ids_tensor in enumerate(token_ids_tensor_list):
         token_ids_cpu_tensor[i][:token_ids_tensor.size(0)] = token_ids_tensor
 
-    pcp_manager.init_batch_info(np.array(list(num_scheduled_tokens.values())), num_reqs)
+    num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
+    for i, req_id in enumerate(req_ids):
+        if num_computed_tokens[i] > 0:
+            num_prompt_tokens[i] = num_computed_tokens[i]
+        else:
+            num_prompt_tokens[i] = num_scheduled_tokens[req_id]
+    input_batch.num_prompt_tokens = num_prompt_tokens
+
+    pcp_manager.init_batch_info(
+        np.array(list(num_scheduled_tokens.values())),
+        num_reqs,
+        input_batch.num_computed_tokens_cpu,
+        input_batch.num_prompt_tokens,
+    )
     pcp_manager.generate_pcp_mtp_input(total_num_scheduled_tokens, num_scheduled_tokens, False,
                                        input_batch, arange_np)
     assert torch.equal(

vllm_ascend/_310p/model_runner_310p.py

@@ -290,6 +290,8 @@ def _prepare_inputs(  # type: ignore[override]
             self.pcp_manager.init_batch_info(
                 num_scheduled_tokens,
                 self.input_batch.num_reqs,
+                self.input_batch.num_computed_tokens_cpu,
+                self.input_batch.num_prompt_tokens,
             )
 
         if self.speculative_config and self.use_cp:

vllm_ascend/attention/context_parallel/attention_cp.py

@@ -116,7 +116,9 @@ def build(
         query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1]
 
         num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = split_decodes_and_prefills(
-            common_attn_metadata, decode_threshold=self.decode_threshold
+            common_attn_metadata,
+            decode_threshold=self.decode_threshold,
+            treat_short_extends_as_decodes=False,
         )
         assert num_decodes + num_prefills == num_reqs
         assert num_decode_tokens + num_prefill_tokens == num_actual_tokens

vllm_ascend/attention/context_parallel/dsa_cp.py

@@ -270,7 +270,11 @@ def build(
         num_input_tokens = common_attn_metadata.num_input_tokens
         if self.common_ratio_to_sas_metadata.get("input_positions", None) is None:
             self.num_decodes, self.num_prefills, self.num_decode_tokens, self.num_prefill_tokens = (
-                split_decodes_and_prefills(common_attn_metadata, decode_threshold=self.decode_threshold)
+                split_decodes_and_prefills(
+                    common_attn_metadata,
+                    decode_threshold=self.decode_threshold,
+                    treat_short_extends_as_decodes=False,
+                )
             )
             self.common_ratio_to_sas_metadata["num_decodes"] = self.num_decodes
             self.common_ratio_to_sas_metadata["num_prefills"] = self.num_prefills
@@ -346,7 +350,9 @@ def build_for_drafting(
         num_reqs = common_attn_metadata.num_reqs
         num_input_tokens = common_attn_metadata.num_input_tokens
         num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(
-            common_attn_metadata, decode_threshold=self.decode_threshold
+            common_attn_metadata,
+            decode_threshold=self.decode_threshold,
+            treat_short_extends_as_decodes=False,
         )
 
         self.num_decodes = num_decodes

vllm_ascend/attention/context_parallel/sfa_cp.py

@@ -93,7 +93,9 @@ def build(
     ) -> AscendSFAMetadata:
         metadata_cls = super().build(common_prefix_len, common_attn_metadata, fast_build)
         num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = split_decodes_and_prefills(
-            common_attn_metadata, decode_threshold=self.decode_threshold
+            common_attn_metadata,
+            decode_threshold=self.decode_threshold,
+            treat_short_extends_as_decodes=False,
         )
         num_reqs = common_attn_metadata.num_reqs
         assert num_decodes + num_prefills == num_reqs

vllm_ascend/attention/mla_v1.py

@@ -437,7 +437,11 @@ def build(
         query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
 
         self.num_decodes, self.num_prefills, self.num_decode_tokens, self.num_prefill_tokens = (
-            split_decodes_and_prefills(common_attn_metadata, decode_threshold=self.decode_threshold)
+            split_decodes_and_prefills(
+                common_attn_metadata,
+                decode_threshold=self.decode_threshold,
+                treat_short_extends_as_decodes=common_attn_metadata.prefill_context_parallel_metadata is None,
+            )
         )
         self.set_num_actual_tokens(common_attn_metadata)
         assert self.num_decodes + self.num_prefills == num_reqs

vllm_ascend/attention/utils.py

@@ -273,17 +273,29 @@ def filter_chunked_req_indices(
 def split_decodes_and_prefills(
     common_attn_metadata: AscendCommonAttentionMetadata,
     decode_threshold: int = 1,
+    require_uniform: bool = False,
+    treat_short_extends_as_decodes: bool = True,
 ) -> tuple[int, int, int, int]:
     """
     Assuming a reordered batch, finds the boundary between prefill and decode
     requests.
     While pcp > 1, query_lens is split across pcp ranks, so we pass in the
     original query_lens and max_query_len to distinguish prefills and decodes.
 
+    The batch is expected to be ordered as:
+    decode -> short_extend -> long_extend -> prefill
+
     Args:
         common_attn_metadata: AscendCommonAttentionMetadata object containing the
             batch metadata.
         decode_threshold: The maximum query length to be considered a decode.
+        require_uniform: If True, requires that all decode requests have the
+            same query length. When set, some queries may be considered
+            prefills even if they are <= decode_threshold, in order to ensure
+            uniformity.
+        treat_short_extends_as_decodes: If True (default), short extends
+            (query_len <= threshold but still prefilling) are counted as
+            decodes. If False, they are counted as prefills.
 
     Returns:
         num_decodes: The number of decode requests.
@@ -296,14 +308,43 @@ def split_decodes_and_prefills(
     max_query_len_pcp_full = long_seq_metadata.max_query_len_pcp_full if long_seq_metadata else 0
     max_query_len = common_attn_metadata.max_query_len if max_query_len_pcp_full == 0 else max_query_len_pcp_full
     num_reqs = common_attn_metadata.num_reqs
+    if num_reqs == 0:
+        return 0, 0, 0, 0
+
     num_tokens = common_attn_metadata.num_actual_tokens
     query_start_loc = common_attn_metadata.query_start_loc_cpu
 
-    if max_query_len <= decode_threshold:
+    if (
+        max_query_len <= decode_threshold
+        and (not require_uniform or decode_threshold <= 1)
+        and treat_short_extends_as_decodes
+    ):
         return num_reqs, 0, num_tokens, 0
 
-    query_lens = (query_start_loc[1:] - query_start_loc[:-1]) if query_lens_pcp_full is None else query_lens_pcp_full
-    is_prefill = query_lens > decode_threshold
+    query_lens_sharded = query_start_loc[1:] - query_start_loc[:-1]
+    query_lens = query_lens_sharded if query_lens_pcp_full is None else query_lens_pcp_full
+    if query_lens[0].item() > decode_threshold:
+        return 0, num_reqs, 0, num_tokens
+
+    if require_uniform:
+        if torch.all((query_lens == query_lens[0]) | (query_lens == 0)):
+            return num_reqs, 0, num_tokens, 0
+        is_prefill = query_lens != query_lens[0]
+    else:
+        is_prefill = query_lens > decode_threshold
+
+    if not treat_short_extends_as_decodes:
+        assert common_attn_metadata.is_prefilling is not None
+        raw_is_prefilling = common_attn_metadata.is_prefilling
+        is_prefilling = raw_is_prefilling[: query_lens.shape[0]]
+        if is_prefilling.shape[0] < query_lens.shape[0]:
+            is_prefilling = F.pad(
+                is_prefilling,
+                (0, query_lens.shape[0] - is_prefilling.shape[0]),
+                value=False,
+            )
+        is_prefill |= is_prefilling
+
     if not torch.any(is_prefill):
         return num_reqs, 0, num_tokens, 0
 

vllm_ascend/ops/gdn_attn_builder.py

@@ -983,6 +983,7 @@ def build(  # type: ignore[override]
             num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = split_decodes_and_prefills(
                 m,
                 decode_threshold=1,
+                treat_short_extends_as_decodes=False,
             )
             num_spec_decode_tokens = 0
             spec_token_indx = None