[BugFix]fix qwen3.5+pcp+chunkprefill accuracy error (vllm-project#11508)

### What this PR does / why we need it?
Fix the precision issue in qwen3.5 with pcp and chunkprefill stacking.

Original code:
updated_state[i] = all_final_state[i] + matmul(all_final_h_update[i],
updated_state[i-1])
Missing - s0
Expanding all_final_state[i] = Φ_i·s0 + p_i:
Original code = (Φ_i·s0 + p_i) + Φ_i·correct_{i-1}
= Φ_i·s0 + p_i + Φ_i·correct_{i-1}
Correct value:
correct_i = p_i + Φ_i·correct_{i-1}
The error is Φ_i·s0—s0's contribution is counted twice (once in
all_final_state[i], and once indirectly through Φ_i·updated_state[i-1]).

Why does this only trigger in chunked prefill?
- First chunk: s0 = 0 (new prompt) → - s0 = - 0 = no-op → the original
code happens to be correct.
- Subsequent chunks: s0 ≠ 0 (carrying state from the previous chunk) →
Φ_i·s0 is recalculated → ssm_state writeback error → precision
divergence.
Thus, short prompts with a single chunk (in memory, non-MTP PCP=2
"correct") cannot detect this; it only becomes apparent when
--max-num-batched-tokens is small enough to split the prompt into
multiple chunks.

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

### How was this patch tested?

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@1f486d9

---------

Signed-off-by: weiguihua2 <weiguihua2@huawei.com>

Author: weiguihua2

tests/ut/ops/a2/test_gdn_chunk_meta.py

@@ -61,6 +61,9 @@ def __setitem__(self, item, value):
     def __add__(self, other):
         return self
 
+    def __sub__(self, other):
+        return self
+
     def transpose(self, dim0, dim1):
         return self
 
@@ -398,3 +401,100 @@ def test_build_final_chunk_indices_falls_back_without_triton_kernel(
     )
 
     assert torch.equal(out_final_chunk_indices, torch.tensor([2, 3, 7], dtype=torch.int32))
+
+
+def test_chunk_gated_delta_rule_fwd_pcp_chaining_subtracts_initial_state(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """PCP chaining uses (updated_state[i-1] - initial_state), not updated_state[i-1].
+
+    With s0 != 0 (subsequent prefill chunk), the fix subtracts s0 to avoid
+    double-counting Φ_i·s0. Verified by checking the returned final_state
+    matches the sequential result Φ_1·(Φ_0·s0+p_0)+p_1.
+    """
+    torch.manual_seed(42)
+    N, H, K, V = 1, 2, 4, 4
+    s0 = torch.randn(N, H, K, V)
+    phi_0 = torch.randn(N, H, K, K)
+    phi_1 = torch.randn(N, H, K, K)
+    p_0 = torch.randn(N, H, K, V)
+    p_1 = torch.randn(N, H, K, V)
+
+    # Each rank computes final_state = Φ_i · s0 + p_i (from shared s0)
+    rank0_fs = torch.matmul(phi_0, s0) + p_0
+    rank1_fs = torch.matmul(phi_1, s0) + p_1
+    # h_update shape [1, N, H, K, K]; after [:, [0], :, :, :] → [1, N, H, K, K]
+    h_update_tensor = phi_0.unsqueeze(0)
+
+    prebuilt_meta = type(
+        "PrebuiltMeta",
+        (),
+        {
+            "block_indices_cumsum": None,
+            "cu_seqlens_host": (0, N),
+            "chunk_indices_chunk64_host": (0, 0),
+            "chunk_indices_chunk64": None,
+            "chunk_offsets_chunk64": torch.tensor([0, 1], dtype=torch.int32),
+            "update_chunk_offsets_chunk64": torch.tensor([0, 2], dtype=torch.int32),
+            "final_chunk_indices_chunk64": torch.tensor([0], dtype=torch.int32),
+            "chunk_indices_large_block": None,
+            "num_decodes": 0,
+        },
+    )()
+
+    all_gather_returns = [
+        torch.stack([rank0_fs, rank1_fs]),  # all_final_state: [2, N, H, K, V]
+        torch.stack([phi_0, phi_1]),  # all_final_h_update: [2, N, H, K, K]
+    ]
+
+    group = type(
+        "Group",
+        (),
+        {
+            "world_size": 2,
+            "rank_in_group": 0,
+            "all_gather": lambda self, value, dim: all_gather_returns.pop(0),
+        },
+    )()
+
+    monkeypatch.setattr(chunk, "get_forward_context", lambda: type("Ctx", (), {"attn_metadata": None})())
+    monkeypatch.setattr(chunk, "get_pcp_group", lambda: group)
+    monkeypatch.setattr(chunk, "chunk_local_cumsum", lambda *a, **kw: _DummyTensor("g_cumsum"))
+    monkeypatch.setattr(chunk, "chunk_scaled_dot_kkt_fwd", lambda *a, **kw: _DummyTensor("A"))
+    monkeypatch.setattr(chunk, "solve_tril", lambda *a, **kw: _DummyTensor("A_solved"))
+    monkeypatch.setattr(chunk, "recompute_w_u_fwd", lambda *a, **kw: (_DummyTensor("w"), _DummyTensor("u")))
+    monkeypatch.setattr(
+        torch.ops._C_ascend,
+        "chunk_gated_delta_rule_fwd_h",
+        lambda *a, **kw: (_DummyTensor("h"), _DummyTensor("v_new"), rank0_fs),
+        raising=False,
+    )
+    monkeypatch.setattr(
+        chunk,
+        "chunk_gated_delta_rule_fwd_hupdate",
+        lambda *a, **kw: h_update_tensor,
+    )
+    monkeypatch.setattr(
+        torch.ops._C_ascend,
+        "chunk_fwd_o",
+        lambda *a, **kw: _DummyTensor("o_ascendc"),
+        raising=False,
+    )
+
+    result = chunk.chunk_gated_delta_rule_fwd(
+        q=_DummyTensor("q"),
+        k=_DummyTensor("k"),
+        v=_DummyTensor("v"),
+        g=_DummyTensor("g"),
+        beta=_DummyTensor("beta"),
+        scale=1.0,
+        initial_state=s0,
+        output_final_state=False,
+        cu_seqlens=torch.tensor([0, N], dtype=torch.int32),
+        prebuilt_meta=prebuilt_meta,
+    )
+
+    final_state = result[3]
+    # Sequential: Φ_1·(Φ_0·s0 + p_0) + p_1
+    expected = torch.matmul(phi_1, torch.matmul(phi_0, s0) + p_0) + p_1
+    torch.testing.assert_close(final_state, expected, rtol=1e-4, atol=1e-4)

vllm_ascend/ops/triton/fla/chunk.py

@@ -142,8 +142,9 @@ def chunk_gated_delta_rule_fwd(
         updated_state = final_state.new_empty(get_pcp_group().world_size, *final_state.shape)
         updated_state[0, ...] = all_final_state[0]
         for i in range(1, get_pcp_group().world_size):
+            # correct_i = all_final_state[i] + Φ_i · (correct_{i-1} - s0) = Φ_i · correct_{i-1} + p_i
             updated_final_state = all_final_state[i] + torch.matmul(
-                all_final_h_update[i, ...], updated_state[i - 1, ...]
+                all_final_h_update[i, ...], updated_state[i - 1, ...] - initial_state
             )
             updated_state[i, ...] = updated_final_state