[Bugfix][DCP] Fix NVIDIA DeepSeek-V3.2 / GLM-5.2 fused attention

tests/kernels/test_fused_deepseek_v32_norm_rope.py

@@ -234,6 +234,63 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
     assert (topk == -1).all(), "topk buffer not cleared on indexer layer"
 
 
+def test_fused_norm_rope_normalizes_query_without_local_cache_slots():
+    """DCP non-owner ranks still need valid query shards for query AllGather."""
+    torch.manual_seed(7)
+    dev = "cuda"
+    num_tokens = 4
+    max_pos = 16
+    pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
+
+    q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
+    kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
+    qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
+    kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
+    ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
+    ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
+    ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
+    cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
+
+    mla_cache = torch.zeros(
+        1, max_pos, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16
+    )
+    idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
+    idx_cache = torch.zeros(1, max_pos, idx_row, device=dev, dtype=torch.uint8)
+    no_local_slots = torch.full((num_tokens,), -1, device=dev, dtype=torch.int64)
+    topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
+
+    q_out = K.fused_norm_rope(
+        pos,
+        q_c,
+        qw,
+        EPS,
+        kv_c,
+        kvw,
+        EPS,
+        k_pe,
+        cos_sin,
+        ik,
+        ikw,
+        ikb,
+        EPS,
+        cos_sin,
+        topk,
+        slot_mapping=no_local_slots,
+        indexer_k_cache=idx_cache,
+        mla_kv_cache=mla_cache,
+        mla_kv_cache_dtype="auto",
+        mla_k_scale=None,
+        has_indexer=True,
+        index_rope_interleave=True,
+    )
+
+    assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm without local cache slots")
+    assert not mla_cache.any(), "non-owner rank wrote the MLA KV cache"
+    assert not idx_cache.any(), "non-owner rank wrote the indexer KV cache"
+    assert (topk == -1).all(), "topk buffer not cleared on non-owner rank"
+
+
 @pytest.mark.parametrize("num_tokens", [1, 17, 512])
 def test_fused_norm_rope_no_indexer(num_tokens: int):
     """Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched."""

vllm/models/deepseek_v32/attention.py

@@ -521,29 +521,45 @@ def _sparse_indexer_and_attn(
             if isinstance(mqa_q_arg, tuple):
                 mqa_q_arg = torch.cat(mqa_q_arg, dim=-1)
             mqa_q_arg = get_tp_group().all_gather(mqa_q_arg, dim=1)
+        elif not self.use_pcp and self.impl.dcp_world_size > 1:
+            assert self.dcp_manager is not None
+            if isinstance(mqa_q_arg, tuple):
+                mqa_q_arg = torch.cat(mqa_q_arg, dim=-1)
+            assert self.dcp_manager.query_gather is not None
+            mqa_q_arg = self.dcp_manager.query_gather(mqa_q_arg)
         attn_out, lse = self.impl.forward_mqa(  # type: ignore[attr-defined]
             mqa_q_arg, kv_cache, attn_metadata, self
         )
 
-        if self.use_pcp and self.impl.dcp_world_size > 1:
+        if self.impl.dcp_world_size > 1:
             assert lse is not None and self.dcp_manager is not None
-            seq_lens = (
-                attn_metadata.decode.seq_lens
-                if attn_metadata.decode is not None
-                else cast(torch.Tensor, attn_metadata.seq_lens)[  # type: ignore[attr-defined]
-                    : attn_metadata.num_decodes
+            seq_lens: torch.Tensor | None
+            query_start_loc: torch.Tensor | None
+            if self.use_pcp:
+                if attn_metadata.decode is not None:
+                    seq_lens = attn_metadata.decode.seq_lens
+                else:
+                    all_seq_lens = cast(
+                        torch.Tensor,
+                        attn_metadata.seq_lens,  # type: ignore[attr-defined]
+                    )
+                    seq_lens = all_seq_lens[: attn_metadata.num_decodes]
+                query_start_loc = attn_metadata.query_start_loc[
+                    : attn_metadata.num_decodes + 1
                 ]
-            )
-            query_start_loc = attn_metadata.query_start_loc[
-                : attn_metadata.num_decodes + 1
-            ]
+            else:
+                # The backend emits (0, -inf) for empty local shards, so no
+                # PCP-only empty-shard metadata is needed.
+                seq_lens = None
+                query_start_loc = None
             attn_out = self.dcp_manager.combine(
                 attn_out,
                 lse,
-                seq_lens=seq_lens,
-                query_start_loc=query_start_loc,
+                seq_lens=seq_lens,  # type: ignore[arg-type]
+                query_start_loc=query_start_loc,  # type: ignore[arg-type]
             )
-            attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
+            if self.use_pcp:
+                attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
 
         # NOTE(woosuk): While the below does not need to be in the eager region,
         # we put it here to avoid copying the attention output. Move this back to the

vllm/models/deepseek_v32/common/kernels.py

@@ -181,22 +181,29 @@ def _fused_norm_rope_kernel(
             )
         return
 
-    if slot_mapping_ptr is None:
-        if kv_out_ptr is None and kpe_out_ptr is None and index_k_out_ptr is None:
-            return
-    elif tl.load(slot_mapping_ptr + tok_idx) < 0:
-        # Padding
-        return
-
     if pid == 2:
-        # Q RMS norm
+        # Q RMS norm. Runs for every row: under DCP a negative slot only
+        # means another rank owns this token's KV slot, but the query is
+        # still needed on every rank (queries are not sharded), so the
+        # slot-based skip below must not gate it. Padding rows do harmless
+        # row-local extra work (no position load, no cache write).
         q_block = tl.arange(0, Q_BLOCK_SIZE)
         q_mask = q_block < Q_DIM
         q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
         q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
         q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
         tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
-    elif pid == 1:
+        return
+
+    if slot_mapping_ptr is None:
+        if kv_out_ptr is None and kpe_out_ptr is None and index_k_out_ptr is None:
+            return
+    elif tl.load(slot_mapping_ptr + tok_idx) < 0:
+        # Padding, or (under DCP) a token whose KV slot another rank owns:
+        # skip the K-side norms and cache writes.
+        return
+
+    if pid == 1:
         # KV RMS Norm + KV RoPE + MLA concat_and_cache.
         # Merged so the normed kv_c and RoPE'd k_pe can be written
         # to the MLA KV cache directly without a separate kernel.