[Feature] qwen3.5 support pcp+mtp (vllm-project#10487)

### What this PR does / why we need it?
This PR introduces support for hybrid attention PCP (Parallel Context
Processing) by adding linear-splitting of prefill inputs
(`_split_pcp_input_hybrid`) and updating the model runner to handle
hybrid attention logits indices. It also corrects several index offsets
in `pcp_utils.py` (switching between decode requests and decode tokens)
and updates Triton operations for Mamba and FLA chunk-gated delta rules.

However, several critical issues were identified during the review:
- In `causal_conv1d.py`, `num_prefills` is initialized to `0` but never
updated from `attn_metadata`, which breaks the PCP + MTP logic.
- In `model_runner_v1.py`, calculating `num_scheduled_tokens` via slice
differences misses the first request's token count, and an
out-of-bounds/incorrect index access occurs when `num_decode_reqs` is
`0`.
- In `chunk.py`, calling `.item()` on an NPU tensor triggers an
expensive device-to-host synchronization in the hot path.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
No testing details were provided in the patch.


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

---------

Signed-off-by: Your Name <you@example.com>
Co-authored-by: Your Name <you@example.com>

Author: SunnyLee151064

tests/ut/spec_decode/a2/test_eagle_proposer.py

@@ -3825,6 +3825,7 @@ def test_set_inputs_first_pass_pcp_dcp_mixed(self):
         self.runner.input_batch.req_ids = req_ids
         # maybe not reasonable just to run test
         self.runner.logits_indices = torch.arange(12, dtype=torch.int32, device=self.device)
+        self.runner.pcp_manager.pcp_use_hybrid_attn = False
 
         proposer, vllm_config = self._create_proposer(
             method="eagle",
@@ -4205,3 +4206,179 @@ def test_set_inputs_first_pass_draft_model(self):
 
         for attrition in attrs_from_cad:
             assert_attr_equal(attrition, expected_cad, out_cad)
+
+
+def _build_split_pcp_input_hybrid_inputs(req_scheduled_tokens: dict[str, int], hidden_size: int):
+    """Build input_ids and target_hidden_states where row i is filled with i.
+
+    Using the row index as the value makes it easy to assert which original
+    tokens survive the per-rank slice and which slots are padding.
+    """
+    total_tokens = sum(req_scheduled_tokens.values())
+    input_ids = torch.arange(total_tokens, dtype=torch.int32)
+    target_hidden_states = torch.arange(total_tokens, dtype=torch.float32).unsqueeze(-1).repeat(1, hidden_size)
+    return input_ids, target_hidden_states
+
+
+# yapf: disable
+@pytest.mark.parametrize(
+    "pcp_size, pcp_rank, req_scheduled_tokens, hidden_size,"
+    " expected_num_tokens, expected_input_ids, expected_hidden_first_col,"
+    " expected_seq_lens, expected_cu_num_tokens, expected_max_query_len",
+    [
+        # Case 1: single req, perfectly aligned to 2*pcp_size.
+        # ori=8, pcp=2 -> padded=8, pcp_tokens=4
+        # rank 0: [0,1,2,3], rank 1: [4,5,6,7]
+        (
+            2, 0, {"0": 8}, 4,
+            4, [0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0],
+            [4], [0, 4], 4,
+        ),
+        (
+            2, 1, {"0": 8}, 4,
+            4, [4, 5, 6, 7], [4.0, 5.0, 6.0, 7.0],
+            [4], [0, 4], 4,
+        ),
+        # Case 2: single req, needs padding.
+        # ori=7, pcp=2 -> padded=8, pcp_tokens=4, num_pads=1
+        # rank 0: [0,1,2,3] (all valid)
+        # rank 1: [4,5,6,PAD] -> input_ids=[4,5,6,0], hidden_first_col=[4,5,6,0]
+        (
+            2, 0, {"0": 7}, 4,
+            4, [0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0],
+            [4], [0, 4], 4,
+        ),
+        (
+            2, 1, {"0": 7}, 4,
+            4, [4, 5, 6, 0], [4.0, 5.0, 6.0, 0.0],
+            [4], [0, 4], 4,
+        ),
+        # Case 3: multiple reqs with different padding needs.
+        # req 0: ori=4 -> padded=4, pcp_tokens=2
+        # req 1: ori=6 -> padded=8, pcp_tokens=4, num_pads=2
+        # rank 0: [0,1] + [4,5,6,7]
+        # rank 1: [2,3] + [8,9,PAD,PAD]
+        (
+            2, 0, {"0": 4, "1": 6}, 2,
+            6, [0, 1, 4, 5, 6, 7], [0.0, 1.0, 4.0, 5.0, 6.0, 7.0],
+            [2, 4], [0, 2, 6], 4,
+        ),
+        (
+            2, 1, {"0": 4, "1": 6}, 2,
+            6, [2, 3, 8, 9, 0, 0], [2.0, 3.0, 8.0, 9.0, 0.0, 0.0],
+            [2, 4], [0, 2, 6], 4,
+        ),
+        # Case 4: pcp_size=4
+        # ori=9, pcp=4 -> padded=16, pcp_tokens=4, num_pads=7
+        # rank 0: tokens [0,1,2,3]
+        # rank 1: tokens [4,5,6,7]
+        # rank 2: tokens [8,PAD,PAD,PAD]
+        # rank 3: tokens [PAD,PAD,PAD,PAD]
+        (
+            4, 0, {"0": 9}, 2,
+            4, [0, 1, 2, 3], [0.0, 1.0, 2.0, 3.0],
+            [4], [0, 4], 4,
+        ),
+        (
+            4, 1, {"0": 9}, 2,
+            4, [4, 5, 6, 7], [4.0, 5.0, 6.0, 7.0],
+            [4], [0, 4], 4,
+        ),
+        (
+            4, 2, {"0": 9}, 2,
+            4, [8, 0, 0, 0], [8.0, 0.0, 0.0, 0.0],
+            [4], [0, 4], 4,
+        ),
+        (
+            4, 3, {"0": 9}, 2,
+            4, [0, 0, 0, 0], [0.0, 0.0, 0.0, 0.0],
+            [4], [0, 4], 4,
+        ),
+        # Case 5: minimal request - single token.
+        # ori=1, pcp=2 -> padded=4, pcp_tokens=2, num_pads=3
+        # rank 0: [0, PAD]
+        # rank 1: [PAD, PAD]
+        (
+            2, 0, {"0": 1}, 2,
+            2, [0, 0], [0.0, 0.0],
+            [2], [0, 2], 2,
+        ),
+        (
+            2, 1, {"0": 1}, 2,
+            2, [0, 0], [0.0, 0.0],
+            [2], [0, 2], 2,
+        ),
+    ],
+)
+# yapf: enable
+def test_split_pcp_input_hybrid(
+    pcp_size,
+    pcp_rank,
+    req_scheduled_tokens,
+    hidden_size,
+    expected_num_tokens,
+    expected_input_ids,
+    expected_hidden_first_col,
+    expected_seq_lens,
+    expected_cu_num_tokens,
+    expected_max_query_len,
+):
+    input_ids, target_hidden_states = _build_split_pcp_input_hybrid_inputs(
+        req_scheduled_tokens, hidden_size
+    )
+
+    # _split_pcp_input_hybrid only reads self.pcp_size and self.pcp_rank,
+    # so a MagicMock with just those attributes is enough to drive the
+    # unbound method without instantiating the full proposer.
+    mock_self = MagicMock()
+    mock_self.pcp_size = pcp_size
+    mock_self.pcp_rank = pcp_rank
+
+    (
+        num_tokens,
+        out_input_ids,
+        out_hidden_states,
+        max_query_len,
+        seq_lens,
+        cu_num_tokens,
+    ) = llm_base_proposer.AscendSpecDecodeBaseProposer._split_pcp_input_hybrid(
+        mock_self, req_scheduled_tokens, input_ids, target_hidden_states
+    )
+
+    assert num_tokens == expected_num_tokens
+    assert max_query_len == expected_max_query_len
+    assert torch.equal(
+        out_input_ids, torch.tensor(expected_input_ids, dtype=torch.int32)
+    )
+    assert torch.equal(seq_lens, torch.tensor(expected_seq_lens, dtype=torch.int32))
+    assert torch.equal(
+        cu_num_tokens, torch.tensor(expected_cu_num_tokens, dtype=torch.int64)
+    )
+
+    # hidden_states shape: [num_tokens, hidden_size]
+    assert out_hidden_states.shape == (expected_num_tokens, hidden_size)
+    assert torch.equal(
+        out_hidden_states[:, 0],
+        torch.tensor(expected_hidden_first_col, dtype=torch.float32),
+    )
+
+
+def test_split_pcp_input_hybrid_preserves_hidden_size():
+    """Hidden states must come back with the same hidden dim as the input."""
+    hidden_size = 7
+    req_scheduled_tokens = {"0": 6}
+    input_ids, target_hidden_states = _build_split_pcp_input_hybrid_inputs(
+        req_scheduled_tokens, hidden_size
+    )
+
+    mock_self = MagicMock()
+    mock_self.pcp_size = 2
+    mock_self.pcp_rank = 0
+
+    _, _, out_hidden_states, _, _, _ = (
+        llm_base_proposer.AscendSpecDecodeBaseProposer._split_pcp_input_hybrid(
+            mock_self, req_scheduled_tokens, input_ids, target_hidden_states
+        )
+    )
+
+    assert out_hidden_states.shape[1] == hidden_size

tests/ut/worker/test_pcp_manager.py

@@ -11,6 +11,7 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 
+import math
 from unittest.mock import MagicMock
 
 import numpy as np
@@ -424,3 +425,116 @@ def test_generate_pcp_mtp_input(
         target_input_ids_pcp_full)
     assert torch.equal(pcp_manager.query_start_loc_pcp_full.cpu[:num_reqs + 1],
                        target_query_start_loc_pcp_full)
+
+
+def _realistic_num_pcp_pads(
+    num_scheduled_tokens: list[int],
+    pcp_world_size: int,
+    num_decode_reqs: int,
+) -> list[int]:
+    """Compute num_pcp_pads exactly as PCPManager.update_tokens_for_pcp would.
+
+    Decode reqs duplicate tokens across pcp_world_size ranks, so their pad
+    count is num_tokens * (pcp_world_size - 1). Prefill reqs are padded up
+    to a multiple of 2 * pcp_world_size, so their pad count is the difference
+    between the padded and original length.
+    """
+    pads: list[int] = []
+    pad_multiple = 2 * pcp_world_size
+    for i, n in enumerate(num_scheduled_tokens):
+        if i < num_decode_reqs:
+            pads.append(n * (pcp_world_size - 1))
+        else:
+            padded = math.ceil(n / pad_multiple) * pad_multiple
+            pads.append(padded - n)
+    return pads
+
+
+# yapf: disable
+@pytest.mark.parametrize(
+    "pcp_size, num_scheduled_tokens, num_decode_reqs,"
+    " expected_cu_num_scheduled_tokens",
+    [
+        # Case 1: no prefill reqs -> returned unchanged.
+        (2, [1, 1], 2, [1, 2]),
+        # Case 2: prefill only (num_decode_reqs == 0). Exercises the case where
+        # the diff-based per-req length derivation would otherwise drop the
+        # first req's length.
+        #   num_scheduled=[3, 5], realistic pads=[1, 3] -> cumsum=[1, 4]
+        #   cu=[3, 8]; padded prefill_lens=[4, 8]; base=0
+        #   prefill_cu=[4, 12]; final = [4*2-1, 12*2-4] = [7, 20]
+        (2, [3, 5], 0, [7, 20]),
+        # Case 3: mix decode + prefill, pcp_size=2.
+        #   num_scheduled=[1, 1, 3, 5], realistic pads=[1, 1, 1, 3]
+        #   prefill pads cumsum=[1, 4]
+        #   cu=[1, 2, 5, 10]; padded prefill_lens=[4, 8]; base=cu[1]=2
+        #   prefill_cu=[6, 14]; final[2:] = [12, 28] - [1, 4] = [11, 24]
+        (2, [1, 1, 3, 5], 2, [1, 2, 11, 24]),
+        # Case 4: pcp_size=4, mix decode + prefill with uneven prefill tokens.
+        #   num_scheduled=[1, 1, 5, 9], realistic pads=[3, 3, 3, 7]
+        #   prefill pads cumsum=[3, 10]
+        #   cu=[1, 2, 7, 16]; padded prefill_lens=[8, 16]; base=cu[1]=2
+        #   prefill_cu=[10, 26]; final[2:] = [40, 104] - [3, 10] = [37, 94]
+        (4, [1, 1, 5, 9], 2, [1, 2, 37, 94]),
+        # Case 5: single prefill req, pcp_size=2.
+        #   num_scheduled=[7], realistic pads=[1]
+        #   cu=[7]; padded prefill_lens=[8]; base=0
+        #   prefill_cu=[8]; final = [8*2-1] = [15]
+        (2, [7], 0, [15]),
+        # Case 6: prefill req that's already aligned to 2*pcp_size.
+        #   num_scheduled=[8], realistic pads=[0]
+        #   cu=[8]; padded prefill_lens=[8]; base=0
+        #   prefill_cu=[8]; final = [8*2-0] = [16]
+        (2, [8], 0, [16]),
+    ],
+)
+# yapf: enable
+def test_adjust_cu_num_scheduled_tokens_for_pcp(
+    pcp_size,
+    num_scheduled_tokens,
+    num_decode_reqs,
+    expected_cu_num_scheduled_tokens,
+):
+    vllm_config = MagicMock()
+    vllm_config.model_config = MagicMock()
+    vllm_config.speculative_config.num_speculative_tokens = 0
+
+    pcp_manager = PCPManager(
+        pcp_world_size=pcp_size,
+        pcp_rank=0,
+        dcp_world_size=1,
+        dcp_rank=0,
+        max_buffer_num_tokens=10000,
+        max_num_reqs=1000,
+        device="cpu",
+        vllm_config=vllm_config,
+        use_async_scheduling=False,
+        pin_memory=False,
+    )
+
+    num_reqs = len(num_scheduled_tokens)
+    cu_num_scheduled_tokens = np.cumsum(
+        np.array(num_scheduled_tokens, dtype=np.int32)
+    )
+    # Use realistic pads that match what update_tokens_for_pcp would produce
+    # for the given num_scheduled_tokens / pcp_world_size combination.
+    # Tests previously fed zeros for prefill pads, which made the math look
+    # correct but never exercised the real runtime path.
+    num_pcp_pads = np.array(
+        _realistic_num_pcp_pads(num_scheduled_tokens, pcp_size, num_decode_reqs),
+        dtype=np.int32,
+    )
+
+    # Seed the manager state normally populated by init_batch_info.
+    pcp_manager.num_decode_reqs = num_decode_reqs
+    pcp_manager.num_prefill_reqs = num_reqs - num_decode_reqs
+
+    result = pcp_manager.adjust_cu_num_scheduled_tokens_for_pcp(
+        cu_num_scheduled_tokens, num_pcp_pads
+    )
+
+    assert np.array_equal(
+        result, np.array(expected_cu_num_scheduled_tokens, dtype=np.int32)
+    ), (
+        f"Expected {expected_cu_num_scheduled_tokens}, got {result.tolist()}"
+    )

vllm_ascend/attention/context_parallel/attention_cp.py

@@ -252,9 +252,9 @@ def build(
                 num_decodes:
             ]
         else:
-            actual_seq_lengths_q = [self.decode_threshold * (i + 1) for i in range(num_decodes)] + query_start_loc_cpu[
-                num_decodes + 1 :
-            ].tolist()
+            actual_seq_lengths_q = (
+                query_start_loc_cpu[1 : num_decodes + 1].tolist() + query_start_loc_cpu[num_decodes + 1 :].tolist()
+            )
 
         attn_metadata = AscendMetadata(
             num_actual_tokens=num_actual_tokens,
@@ -898,7 +898,9 @@ def _gather_and_restore_pcp_qkv(
         decode_offset = attn_metadata.num_decode_tokens * self.pcp_size
         qkv_fa_padding_workspace[:decode_offset] = actual_qkv[:decode_offset]
 
-        pcp_unpad_mask = attn_metadata.prefill.pcp_metadata.pcp_unpad_mask[attn_metadata.num_decodes * self.pcp_size :]
+        pcp_unpad_mask = attn_metadata.prefill.pcp_metadata.pcp_unpad_mask[
+            attn_metadata.num_decode_tokens * self.pcp_size :
+        ]
         qkv_fa_padding_workspace[decode_offset:][pcp_unpad_mask] = actual_qkv[decode_offset:]
         qkv_fa_padding_workspace[decode_offset:][~pcp_unpad_mask] = 0
 

vllm_ascend/ops/gdn_attn_builder.py

@@ -753,17 +753,18 @@ def __init__(
             self.vllm_config.scheduler_config.max_num_seqs,
             self.decode_cudagraph_max_bs,
         )
+
+        self.spec_sequence_masks: torch.Tensor = torch.empty(
+            (sequence_index_capacity,), dtype=torch.bool, device=device
+        )
+
         self.spec_sequence_masks_cpu: torch.Tensor = torch.empty(
             (sequence_index_capacity,),
             dtype=torch.bool,
             device="cpu",
             pin_memory=device.type != "cpu",
         )
-        self.spec_sequence_masks: torch.Tensor = torch.empty(
-            (sequence_index_capacity,),
-            dtype=torch.bool,
-            device=device,
-        )
+
         self.spec_sequence_indices_cpu: torch.Tensor = torch.empty(
             (sequence_index_capacity,),
             dtype=torch.int64,
@@ -796,7 +797,7 @@ def _init_reorder_batch_threshold(
         super()._init_reorder_batch_threshold(
             reorder_batch_threshold,
             supports_spec_as_decode,
-            supports_dcp_with_varlen,
+            True,
         )
         if self.reorder_batch_threshold != 1:  # type: ignore
             speculative_config = self.vllm_config.speculative_config
@@ -1156,6 +1157,7 @@ def build(  # type: ignore[override]
             and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
         ):
             assert spec_sequence_masks is not None
+            self.spec_state_indices_tensor[batch_size:].fill_(NULL_BLOCK_ID)
             self.spec_state_indices_tensor[:num_spec_decodes].copy_(
                 spec_state_indices_tensor,
                 non_blocking=True,

vllm_ascend/ops/triton/fla/chunk.py

@@ -151,7 +151,12 @@ def chunk_gated_delta_rule_fwd(
 
         if get_pcp_group().rank_in_group > 0:
             rerun_initial_state = initial_state.clone()
-            prefill_slice = slice(num_decodes, final_state.shape[0])
+            if cu_seqlens is not None:
+                _ns_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+                prefill_seq_offset = int(((_ns_lens > 0) & (_ns_lens <= 1)).sum().item())
+            else:
+                prefill_seq_offset = num_decodes
+            prefill_slice = slice(prefill_seq_offset, final_state.shape[0])
             rerun_initial_state[prefill_slice] = updated_h_state[prefill_slice]
             h, v_new, _ = chunk_gated_delta_rule_fwd_h(
                 k=k,