[Bugfix][DCP] Materialize prefill keys on non-owner ranks (#54908)
Signed-off-by: 云挚 <ningyunxiao.nyx@antgroup.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Yunxiao Ning committed
Sep 4, 2026 at 10:12 UTC
19c018ec05a3be91e2246be6b6f98971e00d6a1c
2 files changed
+20
-4
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+11
-1
@@ -234,7 +234,7 @@ 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():
237
+def test_fused_norm_rope_materializes_qk_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"
@@ -259,6 +259,8 @@ def test_fused_norm_rope_normalizes_query_without_local_cache_slots():
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
+ kv_out = torch.empty_like(kv_c)
263
+ kpe_out = torch.empty_like(k_pe)
264
265
q_out = K.fused_norm_rope(
266
pos,
@@ -283,9 +285,17 @@ def test_fused_norm_rope_normalizes_query_without_local_cache_slots():
285
mla_k_scale=None,
286
has_indexer=True,
287
index_rope_interleave=True,
288
+ kv_c_out=kv_out,
289
+ k_pe_out=kpe_out,
290
)
291
292
assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm without local cache slots")
293
+ assert_bf16(kv_out, rms_norm(kv_c, kvw), "kv_c rmsnorm without local slots")
294
+ assert_bf16(
295
+ kpe_out,
296
+ rope(k_pe.float(), pos, cos_sin, interleave=True),
297
+ "k_pe RoPE without local slots",
298
+ )
299
assert not mla_cache.any(), "non-owner rank wrote the MLA KV cache"
300
assert not idx_cache.any(), "non-owner rank wrote the indexer KV cache"
301
assert (topk == -1).all(), "topk buffer not cleared on non-owner rank"
vllm/models/deepseek_v32/common/kernels.py
+9
-3
@@ -198,9 +198,13 @@ def _fused_norm_rope_kernel(
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.
201
+ elif tl.load(slot_mapping_ptr + tok_idx) < 0 and (
202
+ pid != 1 or (kv_out_ptr is None and kpe_out_ptr is None)
203
+ ):
204
+ # Padding, or (under DCP) a token whose KV slot another rank owns.
205
+ # Dense prefill still consumes the materialized normalized/rotated K
206
+ # rows on every rank, so pid 1 must produce them even when this rank
207
+ # must not write the owner-local cache slot.
208
return
209
210
if pid == 1:
@@ -250,6 +254,8 @@ def _fused_norm_rope_kernel(
254
return
255
256
slot_idx = tl.load(slot_mapping_ptr + tok_idx)
257
+ if slot_idx < 0:
258
+ return
259
mla_block_size = MLA_CACHE_BLOCK_SIZE
260
mla_block_idx = slot_idx // mla_block_size
261
mla_block_off = slot_idx % mla_block_size