[Bugfix][Attention] Stabilize sparse-MLA DCP for GLM PCP evals (#55879)
Signed-off-by: khluu <khluu000@gmail.com> Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Jason Yao <wsyjh8@gmail.com>
Kevin H. Luu committed
Sep 15, 2026 at 09:58 UTC
241e9391ed85f6c967432135427717a529776a63
10 files changed
+411
-13
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml
+4
-1
@@ -3,10 +3,13 @@ accuracy_threshold: 0.90
3
num_questions: 1319
4
num_fewshot: 5
5
max_concurrency: 100
6
+# Preserve KV-cache headroom after CUDA graph memory profiling and bound the
7
+# DCP-aware FlashInfer workspace to 2.39 GiB at the 4,096-token scheduler cap.
8
server_args: >-
9
--compilation-config '{"cudagraph_mode":"PIECEWISE"}'
10
+ --gpu-memory-utilization 0.96
11
--max-model-len 4096
9
- --max-num-batched-tokens 32768
12
+ --max-num-batched-tokens 4096
13
--safetensors-load-strategy prefetch
14
--moe-backend flashinfer_cutlass
15
--prefill-context-parallel-size 4
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
+2
@@ -3,8 +3,10 @@ accuracy_threshold: 0.90
3
num_questions: 1319
4
num_fewshot: 5
5
max_concurrency: 100
6
+# Leave headroom for FlashInfer CUTLASS autotuning after KV-cache allocation.
7
server_args: >-
8
--compilation-config '{"cudagraph_mode":"NONE"}'
9
+ --gpu-memory-utilization 0.85
10
--max-model-len 4096
11
--max-num-batched-tokens 32768
12
--safetensors-load-strategy prefetch
tests/v1/attention/test_dcp_a2a_pack_mask.py
+15
-1
@@ -13,7 +13,7 @@ from vllm.v1.attention.ops.dcp import (
13
mask_dcp_empty_shards_,
14
)
15
16
-pytestmark = pytest.mark.skipif(
16
+requires_accelerator = pytest.mark.skipif(
17
not current_platform.is_cuda_alike(),
18
reason="needs a CUDA or ROCm device for the Triton pack kernel",
19
)
@@ -44,6 +44,7 @@ def _pack(out, lse, world_size, h_per_rank, head_dim, seq_lens, query_start_loc)
44
@pytest.mark.parametrize("world_size,h_per_rank", [(8, 2), (8, 16), (4, 4), (2, 1)])
45
@pytest.mark.parametrize("tokens_per_req", [1, 2, 3])
46
@pytest.mark.parametrize("num_pad_rows", [0, 2])
47
+@requires_accelerator
48
@pytest.mark.parametrize(
49
"seq_lens_list",
50
[
@@ -97,6 +98,7 @@ def test_fused_mask_matches_eager(
98
99
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
100
@pytest.mark.parametrize("seed", range(8))
101
+@requires_accelerator
102
def test_fused_mask_matches_eager_ragged(dtype, seed):
103
"""Non-uniform per-request query lengths.
104
@@ -137,6 +139,7 @@ def test_fused_mask_matches_eager_ragged(dtype, seed):
139
)
140
141
142
+@requires_accelerator
143
def test_mask_disabled_is_unmasked():
144
"""Passing no seq_lens/query_start_loc must leave the LSE untouched."""
145
torch.manual_seed(0)
@@ -149,3 +152,14 @@ def test_mask_disabled_is_unmasked():
152
153
packed = _pack(out, lse, world_size, h_per_rank, head_dim, None, None)
154
assert torch.isfinite(packed.float()).all()
155
+
156
+
157
+def test_eager_mask_handles_rank_with_no_local_sequences():
158
+ """Padded graph rows are all empty when a DCP rank has no sequences."""
159
+ lse = torch.randn(8, 4, dtype=torch.float32)
160
+ seq_lens = torch.empty(0, dtype=torch.int32)
161
+ query_start_loc = torch.tensor([0], dtype=torch.int32)
162
+
163
+ mask_dcp_empty_shards_(lse, seq_lens, query_start_loc)
164
+
165
+ assert torch.isneginf(lse).all()
tests/v1/attention/test_flashinfer_sparse_mla_workspace.py
new
+166
@@ -0,0 +1,166 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU-only tests for FlashInfer sparse MLA backend constraints."""
4
+
5
+from types import SimpleNamespace
6
+from unittest.mock import patch
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
12
+ _DEFAULT_WORKSPACE_BUFFER_SIZE,
13
+ FlashInferMLASparseTRTLLMBackend,
14
+ _required_workspace_bytes,
15
+ compute_trtllm_sparse_mla_workspace_bytes,
16
+)
17
+
18
+# Reporter config from vllm#50781: GLM-5.2 (64 q heads), TP=8, DCP=8,
19
+# max_num_batched_tokens=16384, FlashInfer 0.6.14.
20
+_REPORTER_HEADS_PER_RANK = 64 // 8
21
+_REPORTER_DCP = 8
22
+_REPORTER_MAX_TOKENS = 16384
23
+# From the reported crash: trtllm-gen requested exactly this many bytes for
24
+# trtllm_gen_softmax_workspace after an 8 MiB counter carve had left
25
+# 404,750,336 of the 413,138,944-byte default buffer.
26
+_REPORTER_OBSERVED_SOFTMAX_BYTES = 1_611_661_312
27
+_REPORTER_OBSERVED_PRECARVE_BYTES = 8_388_608
28
+
29
+
30
+def test_reporter_softmax_carve_is_reproduced_byte_exactly():
31
+ # The crashing step scheduled 12288 tokens (a single 12K-token request);
32
+ # the softmax slab formula must reproduce the kernel's request exactly:
33
+ # 8 * (8 heads/rank * 8 dcp) * 12288 tokens * 256 + 1 MiB guard.
34
+ softmax_bytes = compute_trtllm_sparse_mla_workspace_bytes(
35
+ base_workspace_bytes=0,
36
+ dcp_world_size=_REPORTER_DCP,
37
+ num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
38
+ max_num_batched_tokens=12288,
39
+ )
40
+ assert softmax_bytes == _REPORTER_OBSERVED_SOFTMAX_BYTES
41
+
42
+
43
+def test_reporter_config_covers_observed_overflow():
44
+ computed = compute_trtllm_sparse_mla_workspace_bytes(
45
+ base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE,
46
+ dcp_world_size=_REPORTER_DCP,
47
+ num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
48
+ max_num_batched_tokens=_REPORTER_MAX_TOKENS,
49
+ )
50
+ assert computed >= (
51
+ _REPORTER_OBSERVED_SOFTMAX_BYTES + _REPORTER_OBSERVED_PRECARVE_BYTES
52
+ )
53
+
54
+
55
+def test_non_dcp_size_is_unchanged():
56
+ computed = compute_trtllm_sparse_mla_workspace_bytes(
57
+ base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE,
58
+ dcp_world_size=1,
59
+ num_heads_per_rank=128,
60
+ max_num_batched_tokens=65536,
61
+ )
62
+ assert computed == _DEFAULT_WORKSPACE_BUFFER_SIZE
63
+
64
+
65
+def test_combined_pcp_dcp_rejects_flashinfer_sparse():
66
+ config = SimpleNamespace(
67
+ parallel_config=SimpleNamespace(
68
+ prefill_context_parallel_size=4,
69
+ decode_context_parallel_size=4,
70
+ ),
71
+ model_config=None,
72
+ )
73
+ with patch("vllm.config.get_current_vllm_config", return_value=config):
74
+ reason = FlashInferMLASparseTRTLLMBackend.supports_combination(
75
+ head_size=576,
76
+ dtype=torch.bfloat16,
77
+ kv_cache_dtype="fp8",
78
+ block_size=64,
79
+ use_mla=True,
80
+ has_sink=False,
81
+ use_sparse=True,
82
+ use_mm_prefix=False,
83
+ device_capability=SimpleNamespace(major=10),
84
+ )
85
+
86
+ assert reason is not None
87
+ assert "use FLASHMLA_SPARSE" in reason
88
+
89
+
90
+@pytest.mark.parametrize("pcp_size,dcp_size", [(4, 1), (1, 4)])
91
+def test_single_context_parallel_mode_allows_flashinfer_sparse(
92
+ pcp_size: int, dcp_size: int
93
+):
94
+ config = SimpleNamespace(
95
+ parallel_config=SimpleNamespace(
96
+ prefill_context_parallel_size=pcp_size,
97
+ decode_context_parallel_size=dcp_size,
98
+ ),
99
+ model_config=None,
100
+ )
101
+ with patch("vllm.config.get_current_vllm_config", return_value=config):
102
+ reason = FlashInferMLASparseTRTLLMBackend.supports_combination(
103
+ head_size=576,
104
+ dtype=torch.bfloat16,
105
+ kv_cache_dtype="fp8",
106
+ block_size=64,
107
+ use_mla=True,
108
+ has_sink=False,
109
+ use_sparse=True,
110
+ use_mm_prefix=False,
111
+ device_capability=SimpleNamespace(major=10),
112
+ )
113
+
114
+ assert reason is None
115
+
116
+
117
+def test_default_constant_matches_envs_default(monkeypatch):
118
+ from vllm import envs
119
+
120
+ monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False)
121
+ assert _DEFAULT_WORKSPACE_BUFFER_SIZE == envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
122
+
123
+
124
+def test_env_unset_returns_computed(monkeypatch):
125
+ monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False)
126
+ required = _required_workspace_bytes(
127
+ dcp_world_size=_REPORTER_DCP,
128
+ num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
129
+ max_num_batched_tokens=_REPORTER_MAX_TOKENS,
130
+ )
131
+ assert required == compute_trtllm_sparse_mla_workspace_bytes(
132
+ _DEFAULT_WORKSPACE_BUFFER_SIZE,
133
+ _REPORTER_DCP,
134
+ _REPORTER_HEADS_PER_RANK,
135
+ _REPORTER_MAX_TOKENS,
136
+ )
137
+
138
+
139
+def test_env_override_below_computed_is_respected_with_warning(monkeypatch):
140
+ override = 100 * 1024 * 1024
141
+ monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override))
142
+ with patch(
143
+ "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger"
144
+ ) as mock_logger:
145
+ required = _required_workspace_bytes(
146
+ dcp_world_size=_REPORTER_DCP,
147
+ num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
148
+ max_num_batched_tokens=_REPORTER_MAX_TOKENS,
149
+ )
150
+ assert required == override
151
+ mock_logger.warning_once.assert_called_once()
152
+
153
+
154
+def test_env_override_above_computed_is_respected_without_warning(monkeypatch):
155
+ override = 8 * 1024 * 1024 * 1024
156
+ monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override))
157
+ with patch(
158
+ "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger"
159
+ ) as mock_logger:
160
+ required = _required_workspace_bytes(
161
+ dcp_world_size=_REPORTER_DCP,
162
+ num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
163
+ max_num_batched_tokens=_REPORTER_MAX_TOKENS,
164
+ )
165
+ assert required == override
166
+ mock_logger.warning_once.assert_not_called()
tests/v1/attention/test_indexer_dcp_localize.py
+17
@@ -307,6 +307,23 @@ def test_get_dcp_local_seq_lens_preserves_mtp_bounds_shape():
307
torch.testing.assert_close(actual, expected)
308
309
310
+def test_get_dcp_local_seq_lens_reuses_device_rank_tensor(monkeypatch):
311
+ seq_lens = torch.tensor([0, 1, 7, 8, 17], dtype=torch.int32)
312
+ world = 4
313
+ rank = 2
314
+ rank_tensor = torch.tensor(rank, dtype=torch.int32)
315
+ expected = get_dcp_local_seq_lens(seq_lens, world, rank, 2)
316
+
317
+ def fail_tensor_construction(*args, **kwargs):
318
+ raise AssertionError("rank tensor must be reused")
319
+
320
+ with monkeypatch.context() as patch:
321
+ patch.setattr(torch, "tensor", fail_tensor_construction)
322
+ actual = get_dcp_local_seq_lens(seq_lens, world, rank_tensor, 2)
323
+
324
+ torch.testing.assert_close(actual, expected)
325
+
326
+
327
def test_get_dcp_local_seq_lens_must_run_after_decode_expansion():
328
world = 2
329
rank = 1
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+32
@@ -179,6 +179,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
179
),
180
parallel_config=SimpleNamespace(
181
decode_context_parallel_size=1,
182
+ prefill_context_parallel_size=1,
183
cp_kv_cache_interleave_size=1,
184
),
185
)
@@ -193,6 +194,37 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
194
}
195
196
197
+def test_indexer_warmup_includes_pcp_normalized_dcp_key(monkeypatch):
198
+ monkeypatch.setattr(
199
+ "vllm.v1.attention.backends.mla.indexer.get_dcp_group",
200
+ lambda: SimpleNamespace(rank_in_group=2),
201
+ )
202
+ config = SimpleNamespace(
203
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
204
+ model_config=SimpleNamespace(
205
+ hf_text_config=SimpleNamespace(compress_ratios=[32], index_kpool=32)
206
+ ),
207
+ parallel_config=SimpleNamespace(
208
+ decode_context_parallel_size=4,
209
+ prefill_context_parallel_size=4,
210
+ cp_kv_cache_interleave_size=1,
211
+ ),
212
+ )
213
+
214
+ keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config)
215
+
216
+ # Triton's compile key normalizes generic i32 values to 2 and divisible
217
+ # i32 values (including zero) to 16.
218
+ assert {(key.dcp_rank, key.dcp_world) for key in keys} == {(2, 2), (16, 1)}
219
+ assert {
220
+ (
221
+ key.input_variant.is_aligned("uncompressed_seq_lens"),
222
+ key.input_variant.is_aligned("cu_compressed_seq_lens"),
223
+ )
224
+ for key in keys
225
+ } == {(False, False), (False, True), (True, False), (True, True)}
226
+
227
+
228
def test_compressed_slot_mapping_warmup_includes_index_kpool():
229
config = SimpleNamespace(
230
cache_config=SimpleNamespace(block_size=256),
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+122
-3
@@ -113,6 +113,16 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
113
from vllm.config import get_current_vllm_config
114
115
vllm_config = get_current_vllm_config()
116
+ parallel_config = vllm_config.parallel_config
117
+ if (
118
+ parallel_config.prefill_context_parallel_size > 1
119
+ and parallel_config.decode_context_parallel_size > 1
120
+ ):
121
+ return (
122
+ "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; "
123
+ "use FLASHMLA_SPARSE, which gathers each DCP KV shard before "
124
+ "running the rank-local PCP prefill queries"
125
+ )
126
if kv_cache_dtype == "fp8_ds_mla":
127
return (
128
"FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype"
@@ -265,6 +275,31 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde
275
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
276
hisparse_supports_multi_token_decode: ClassVar[bool] = True
277
278
+ def __init__(
279
+ self,
280
+ kv_cache_spec: AttentionSpec,
281
+ layer_names: list[str],
282
+ vllm_config: VllmConfig,
283
+ device: torch.device,
284
+ ) -> None:
285
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
286
+
287
+ # Under DCP the workspace must hold the trtllm-gen softmax-stats slab.
288
+ # The buffer address is baked into CUDA graphs, so it has to reach its
289
+ # final size here, before the first forward/capture (no lazy regrow).
290
+ if self.dcp_world_size > 1:
291
+ num_q_heads = vllm_config.model_config.get_num_attention_heads(
292
+ vllm_config.parallel_config
293
+ )
294
+ _get_workspace_buffer(
295
+ device,
296
+ _required_workspace_bytes(
297
+ self.dcp_world_size,
298
+ num_q_heads,
299
+ vllm_config.scheduler_config.max_num_batched_tokens,
300
+ ),
301
+ )
302
+
303
def _build_req_id_per_token(
304
self,
305
common_attn_metadata: "CommonAttentionMetadata",
@@ -275,14 +310,98 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde
310
# Global workspace buffer (lazily initialized)
311
_fi_sparse_workspace: torch.Tensor | None = None
312
313
+# trtllm-gen carves a softmax-stats slab from the workspace whenever LSE is
314
+# requested (FlashInfer csrc/trtllm_fmha_kernel_launcher.cu, unchanged from
315
+# v0.6.14 through current main):
316
+# sizeof(float2) * num_qo_heads * batch_size * round_up(max_q_len, 256)
317
+# + 1 MiB guard
318
+# forward_mqa always passes q_len == 1, so each (head, token) pair costs
319
+# round_up(1, 256) == 256 slots.
320
+_TRTLLM_GEN_SOFTMAX_STAT_BYTES = 8 # sizeof(float2)
321
+_TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN = 256 # round_up(q_len=1, 256)
322
+_TRTLLM_GEN_SOFTMAX_GUARD_BYTES = 1024 * 1024
323
+
324
+# Keep in sync with the VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE default in
325
+# vllm/envs.py.
326
+_DEFAULT_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024
327
+
328
+
329
+def compute_trtllm_sparse_mla_workspace_bytes(
330
+ base_workspace_bytes: int,
331
+ dcp_world_size: int,
332
+ num_heads_per_rank: int,
333
+ max_num_batched_tokens: int,
334
+) -> int:
335
+ """Workspace bytes needed by the trtllm-gen sparse MLA decode kernel.
336
+
337
+ Under DCP the query is all-gathered across the DCP group in the head dim
338
+ and the kernel is asked for LSE, so trtllm-gen carves the softmax-stats
339
+ slab described above from the workspace before its (batch-independent)
340
+ counter and scratch regions. ``base_workspace_bytes`` must stay available
341
+ for those regions, so the slab is added on top of it (see #50781).
342
+
343
+ Without DCP no LSE is requested, no slab is carved, and the base size is
344
+ returned unchanged.
345
+ """
346
+ if dcp_world_size <= 1:
347
+ return base_workspace_bytes
348
+ softmax_bytes = (
349
+ _TRTLLM_GEN_SOFTMAX_STAT_BYTES
350
+ * (num_heads_per_rank * dcp_world_size)
351
+ * max_num_batched_tokens
352
+ * _TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN
353
+ + _TRTLLM_GEN_SOFTMAX_GUARD_BYTES
354
+ )
355
+ return base_workspace_bytes + softmax_bytes
356
+
357
+
358
+def _required_workspace_bytes(
359
+ dcp_world_size: int,
360
+ num_heads_per_rank: int,
361
+ max_num_batched_tokens: int,
362
+) -> int:
363
+ """Resolve the workspace size, honoring an explicit env override."""
364
+ computed = compute_trtllm_sparse_mla_workspace_bytes(
365
+ _DEFAULT_WORKSPACE_BUFFER_SIZE,
366
+ dcp_world_size,
367
+ num_heads_per_rank,
368
+ max_num_batched_tokens,
369
+ )
370
+ if not envs.is_set("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE"):
371
+ return computed
372
+ env_bytes = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
373
+ if env_bytes < computed:
374
+ logger.warning_once(
375
+ "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d is below the %d bytes "
376
+ "computed for the sparse MLA workspace with dcp_world_size=%d, "
377
+ "%d heads/rank and max_num_batched_tokens=%d. Respecting the "
378
+ "override, but the trtllm-gen kernel may crash with a workspace "
379
+ "overflow; set VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d or unset "
380
+ "it to use the computed size.",
381
+ env_bytes,
382
+ computed,
383
+ dcp_world_size,
384
+ num_heads_per_rank,
385
+ max_num_batched_tokens,
386
+ computed,
387
+ )
388
+ return env_bytes
389
+
390
279
-def _get_workspace_buffer(device: torch.device) -> torch.Tensor:
391
+def _get_workspace_buffer(
392
+ device: torch.device, min_bytes: int | None = None
393
+) -> torch.Tensor:
394
global _fi_sparse_workspace
281
- if _fi_sparse_workspace is None:
395
+ required = (
396
+ min_bytes
397
+ if min_bytes is not None
398
+ else envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
399
+ )
400
+ if _fi_sparse_workspace is None or _fi_sparse_workspace.numel() < required:
401
# FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace;
402
# the trtllm-gen path views it as uint8, so int8 is safe for all backends.
403
_fi_sparse_workspace = torch.zeros(
285
- envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
404
+ required,
405
dtype=torch.int8,
406
device=device,
407
)
vllm/v1/attention/backends/mla/indexer.py
+39
-7
@@ -10,7 +10,7 @@ import vllm.envs as envs
10
from vllm.config import VllmConfig
11
from vllm.distributed import get_dcp_group, get_pcp_group
12
from vllm.logger import init_logger
13
-from vllm.model_executor.warmup.jit_warmup import kernel_launcher
13
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
14
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
15
LaunchSpec,
16
TritonPointerInputVariant,
@@ -493,6 +493,14 @@ class BuildPrefillChunkMetadataKernel(
493
dcp_world = parallel_config.decode_context_parallel_size
494
dcp_interleave = parallel_config.cp_kv_cache_interleave_size
495
dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0
496
+ dcp_cases = [dict(DCP_RANK=dcp_rank, DCP_WORLD=dcp_world)]
497
+ if parallel_config.prefill_context_parallel_size > 1 and dcp_world > 1:
498
+ # PCP+DCP prefill constructs a global PCP chunk plan and therefore
499
+ # dispatches this kernel with already-localized row starts. The
500
+ # runtime key is normalized to rank 0/world 1 in that path; warm it
501
+ # explicitly so one PCP rank cannot enter a collective while
502
+ # another rank is still compiling the first real inference batch.
503
+ dcp_cases.append(dict(DCP_RANK=0, DCP_WORLD=1))
504
compress_ratios = tuple(
505
dict.fromkeys(
506
max(1, int(ratio))
@@ -506,17 +514,33 @@ class BuildPrefillChunkMetadataKernel(
514
if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios:
515
compress_ratios = compress_ratios + (index_kpool,)
516
return self._trace_dispatch(self.dispatch)(
517
+ zip_inputs(*dcp_cases),
518
# Cover Triton's divisible, exact-one, and generic i32 classes.
519
query_slice_start=(0, 1, 2),
520
query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
512
- DCP_RANK=dcp_rank,
513
- DCP_WORLD=dcp_world,
521
DCP_INTERLEAVE=dcp_interleave,
522
BLOCK_SIZE=self.BLOCK_SIZE,
523
COMPRESS_RATIO=list(compress_ratios),
524
+ # PCP's global cumulative lengths are the second row of one packed
525
+ # allocation, so their pointer is not always 16-byte aligned.
526
+ # Prefill chunking can independently unalign uncompressed lengths.
527
input_variant=(
518
- TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
519
- TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
528
+ TritonPointerInputVariant.from_alignment(
529
+ uncompressed_seq_lens=True,
530
+ cu_compressed_seq_lens=True,
531
+ ),
532
+ TritonPointerInputVariant.from_alignment(
533
+ uncompressed_seq_lens=True,
534
+ cu_compressed_seq_lens=False,
535
+ ),
536
+ TritonPointerInputVariant.from_alignment(
537
+ uncompressed_seq_lens=False,
538
+ cu_compressed_seq_lens=True,
539
+ ),
540
+ TritonPointerInputVariant.from_alignment(
541
+ uncompressed_seq_lens=False,
542
+ cu_compressed_seq_lens=False,
543
+ ),
544
),
545
)
546
@@ -527,7 +551,9 @@ class BuildPrefillChunkMetadataKernel(
551
uncompressed_seq_lens=compile_key.input_variant.pointer(
552
"uncompressed_seq_lens", torch.int32
553
),
530
- cu_compressed_seq_lens=int32_ptr,
554
+ cu_compressed_seq_lens=compile_key.input_variant.pointer(
555
+ "cu_compressed_seq_lens", torch.int32
556
+ ),
557
row_start_cu_compressed_seq_lens=int32_ptr,
558
token_to_seq=int32_ptr,
559
cu_compressed_seq_len_ks=int32_ptr,
@@ -836,6 +862,12 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
862
dtype=torch.int32,
863
device=self.device,
864
)
865
+ # Materialize the rank on device during builder initialization. Creating
866
+ # this scalar in build() would introduce a GPU<->CPU sync in the decode
867
+ # hot path.
868
+ self.dcp_rank_tensor = torch.tensor(
869
+ self.dcp_rank, dtype=torch.int32, device=self.device
870
+ )
871
self.expanded_block_table_buffer = torch.zeros(
872
(scheduler_config.max_num_batched_tokens, block_table_width),
873
dtype=torch.int32,
@@ -889,7 +921,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
921
local_seq_lens = get_dcp_local_seq_lens(
922
seq_lens,
923
self.dcp_world_size,
892
- self.dcp_rank,
924
+ self.dcp_rank_tensor,
925
self.cp_kv_cache_interleave_size,
926
)
927
if seq_lens_is_buffer_view:
vllm/v1/attention/backends/utils.py
+7
-1
@@ -1094,7 +1094,7 @@ def compute_causal_conv1d_metadata(
1094
def get_dcp_local_seq_lens(
1095
seq_lens: torch.Tensor,
1096
dcp_size: int = 1,
1097
- dcp_rank: int | None = None,
1097
+ dcp_rank: int | torch.Tensor | None = None,
1098
cp_kv_cache_interleave_size: int = 1,
1099
) -> torch.Tensor:
1100
"""While using dcp, kv_cache size stored on each rank may be different,
@@ -1112,6 +1112,12 @@ def get_dcp_local_seq_lens(
1112
dcp_size,
1113
)
1114
seq_lens_tiled = seq_lens_i32.unsqueeze(-1)
1115
+ elif isinstance(dcp_rank, torch.Tensor):
1116
+ assert dcp_rank.dtype == torch.int32
1117
+ assert dcp_rank.device == seq_lens.device
1118
+ assert dcp_rank.numel() == 1
1119
+ rank_offsets = dcp_rank
1120
+ seq_lens_tiled = seq_lens_i32
1121
else:
1122
rank_offsets = torch.tensor(dcp_rank, dtype=torch.int32, device=seq_lens.device)
1123
seq_lens_tiled = seq_lens_i32
vllm/v1/attention/ops/dcp.py
+7
@@ -77,6 +77,13 @@ def mask_dcp_empty_shards_(
77
return
78
assert seq_lens is not None and query_start_loc is not None
79
80
+ # A DCP rank can receive no local sequences during CUDA graph warmup even
81
+ # though the padded LSE buffer still has rows. In that case every row is an
82
+ # empty shard; avoid indexing the empty seq_lens tensor below.
83
+ if seq_lens.shape[0] == 0:
84
+ lse.fill_(float("-inf"))
85
+ return
86
+
87
row_indices = torch.arange(
88
lse.shape[0], device=lse.device, dtype=query_start_loc.dtype
89
)