[Bugfix][DCP] Fix NVIDIA DeepSeek-V3.2 / GLM-5.2 fused attention (#50005)
Signed-off-by: 云挚 <ningyunxiao.nyx@antgroup.com> Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Yunxiao Ning committed
Aug 31, 2026 at 15:30 UTC
2cf82bcdd17f658486c280b0acd23274f60bdf00
3 files changed
+102
-22
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+57
@@ -234,6 +234,63 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str
234
assert (topk == -1).all(), "topk buffer not cleared on indexer layer"
235
236
237
+def test_fused_norm_rope_normalizes_query_without_local_cache_slots():
238
+ """DCP non-owner ranks still need valid query shards for query AllGather."""
239
+ torch.manual_seed(7)
240
+ dev = "cuda"
241
+ num_tokens = 4
242
+ max_pos = 16
243
+ pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)
244
+
245
+ q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
246
+ kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
247
+ k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
248
+ qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
249
+ kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
250
+ ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)
251
+ ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
252
+ ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)
253
+ cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
254
+
255
+ mla_cache = torch.zeros(
256
+ 1, max_pos, KV_LORA + ROPE_DIM, device=dev, dtype=torch.bfloat16
257
+ )
258
+ idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4
259
+ idx_cache = torch.zeros(1, max_pos, idx_row, device=dev, dtype=torch.uint8)
260
+ no_local_slots = torch.full((num_tokens,), -1, device=dev, dtype=torch.int64)
261
+ topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
262
+
263
+ q_out = K.fused_norm_rope(
264
+ pos,
265
+ q_c,
266
+ qw,
267
+ EPS,
268
+ kv_c,
269
+ kvw,
270
+ EPS,
271
+ k_pe,
272
+ cos_sin,
273
+ ik,
274
+ ikw,
275
+ ikb,
276
+ EPS,
277
+ cos_sin,
278
+ topk,
279
+ slot_mapping=no_local_slots,
280
+ indexer_k_cache=idx_cache,
281
+ mla_kv_cache=mla_cache,
282
+ mla_kv_cache_dtype="auto",
283
+ mla_k_scale=None,
284
+ has_indexer=True,
285
+ index_rope_interleave=True,
286
+ )
287
+
288
+ assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm without local cache slots")
289
+ assert not mla_cache.any(), "non-owner rank wrote the MLA KV cache"
290
+ assert not idx_cache.any(), "non-owner rank wrote the indexer KV cache"
291
+ assert (topk == -1).all(), "topk buffer not cleared on non-owner rank"
292
+
293
+
294
@pytest.mark.parametrize("num_tokens", [1, 17, 512])
295
def test_fused_norm_rope_no_indexer(num_tokens: int):
296
"""Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched."""
vllm/models/deepseek_v32/attention.py
+29
-13
@@ -521,29 +521,45 @@ class DeepseekV32Attention(MLAAttention):
521
if isinstance(mqa_q_arg, tuple):
522
mqa_q_arg = torch.cat(mqa_q_arg, dim=-1)
523
mqa_q_arg = get_tp_group().all_gather(mqa_q_arg, dim=1)
524
+ elif not self.use_pcp and self.impl.dcp_world_size > 1:
525
+ assert self.dcp_manager is not None
526
+ if isinstance(mqa_q_arg, tuple):
527
+ mqa_q_arg = torch.cat(mqa_q_arg, dim=-1)
528
+ assert self.dcp_manager.query_gather is not None
529
+ mqa_q_arg = self.dcp_manager.query_gather(mqa_q_arg)
530
attn_out, lse = self.impl.forward_mqa( # type: ignore[attr-defined]
531
mqa_q_arg, kv_cache, attn_metadata, self
532
)
533
528
- if self.use_pcp and self.impl.dcp_world_size > 1:
534
+ if self.impl.dcp_world_size > 1:
535
assert lse is not None and self.dcp_manager is not None
530
- seq_lens = (
531
- attn_metadata.decode.seq_lens
532
- if attn_metadata.decode is not None
533
- else cast(torch.Tensor, attn_metadata.seq_lens)[ # type: ignore[attr-defined]
534
- : attn_metadata.num_decodes
536
+ seq_lens: torch.Tensor | None
537
+ query_start_loc: torch.Tensor | None
538
+ if self.use_pcp:
539
+ if attn_metadata.decode is not None:
540
+ seq_lens = attn_metadata.decode.seq_lens
541
+ else:
542
+ all_seq_lens = cast(
543
+ torch.Tensor,
544
+ attn_metadata.seq_lens, # type: ignore[attr-defined]
545
+ )
546
+ seq_lens = all_seq_lens[: attn_metadata.num_decodes]
547
+ query_start_loc = attn_metadata.query_start_loc[
548
+ : attn_metadata.num_decodes + 1
549
]
536
- )
537
- query_start_loc = attn_metadata.query_start_loc[
538
- : attn_metadata.num_decodes + 1
539
- ]
550
+ else:
551
+ # The backend emits (0, -inf) for empty local shards, so no
552
+ # PCP-only empty-shard metadata is needed.
553
+ seq_lens = None
554
+ query_start_loc = None
555
attn_out = self.dcp_manager.combine(
556
attn_out,
557
lse,
543
- seq_lens=seq_lens,
544
- query_start_loc=query_start_loc,
558
+ seq_lens=seq_lens, # type: ignore[arg-type]
559
+ query_start_loc=query_start_loc, # type: ignore[arg-type]
560
)
546
- attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
561
+ if self.use_pcp:
562
+ attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
563
564
# NOTE(woosuk): While the below does not need to be in the eager region,
565
# we put it here to avoid copying the attention output. Move this back to the
vllm/models/deepseek_v32/common/kernels.py
+16
-9
@@ -181,22 +181,29 @@ def _fused_norm_rope_kernel(
181
)
182
return
183
184
- if slot_mapping_ptr is None:
185
- if kv_out_ptr is None and kpe_out_ptr is None and index_k_out_ptr is None:
186
- return
187
- elif tl.load(slot_mapping_ptr + tok_idx) < 0:
188
- # Padding
189
- return
190
-
184
if pid == 2:
192
- # Q RMS norm
185
+ # Q RMS norm. Runs for every row: under DCP a negative slot only
186
+ # means another rank owns this token's KV slot, but the query is
187
+ # still needed on every rank (queries are not sharded), so the
188
+ # slot-based skip below must not gate it. Padding rows do harmless
189
+ # row-local extra work (no position load, no cache write).
190
q_block = tl.arange(0, Q_BLOCK_SIZE)
191
q_mask = q_block < Q_DIM
192
q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
193
q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
194
q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
195
tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
199
- elif pid == 1:
196
+ return
197
+
198
+ if slot_mapping_ptr is None:
199
+ if kv_out_ptr is None and kpe_out_ptr is None and index_k_out_ptr is None:
200
+ return
201
+ elif tl.load(slot_mapping_ptr + tok_idx) < 0:
202
+ # Padding, or (under DCP) a token whose KV slot another rank owns:
203
+ # skip the K-side norms and cache writes.
204
+ return
205
+
206
+ if pid == 1:
207
# KV RMS Norm + KV RoPE + MLA concat_and_cache.
208
# Merged so the normed kv_c and RoPE'd k_pe can be written
209
# to the MLA KV cache directly without a separate kernel.