[DCP] Default query replication for GLM sparse attention (#50382)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Lucas Wilkinson committed
Aug 21, 2026 at 12:07 UTC
72aedcc426c23efe01e2484cadb76074330a10a9
6 files changed
+62
-20
tests/distributed/test_dcp_a2a.py
+2
-11
@@ -105,20 +105,11 @@ class TestDCPCommBackendConfig:
105
"""Test --dcp-comm-backend config validation."""
106
107
def test_default_is_ag_rs(self):
108
- """Default comm backend is ag_rs."""
108
+ """Comm backend resolves to ag_rs unless the model asks otherwise."""
109
config = ParallelConfig()
110
+ config.set_dcp_defaults()
111
assert config.dcp_comm_backend == "ag_rs"
112
112
- def test_a2a_requires_dcp_greater_than_1(self):
113
- """A2A backend requires decode_context_parallel_size > 1."""
114
- with pytest.raises(
115
- ValueError, match="requires decode_context_parallel_size > 1"
116
- ):
117
- ParallelConfig(
118
- dcp_comm_backend="a2a",
119
- decode_context_parallel_size=1,
120
- )
121
-
113
def test_a2a_with_dcp_valid(self):
114
"""A2A backend is valid when DCP > 1."""
115
config = ParallelConfig(
vllm/config/parallel.py
+31
-7
@@ -351,12 +351,26 @@ class ParallelConfig:
351
and will be deprecated when PCP is fully supported.
352
353
"""
354
- dcp_comm_backend: DCPCommBackend = "ag_rs"
354
+ dcp_comm_backend: DCPCommBackend | None = None
355
"""Communication backend for Decode Context Parallel (DCP).
356
- - "ag_rs": AllGather + ReduceScatter (default, existing behavior)
356
+ - "ag_rs": AllGather + ReduceScatter (existing behavior)
357
- "a2a": All-to-All exchange of partial outputs + LSE, then
358
combine with Triton kernel. Reduces NCCL calls from 3 to 2
359
per layer for MLA models.
360
+
361
+ `None` selects the model default, which is "ag_rs" unless the model
362
+ overrides it via [`set_dcp_defaults`][vllm.config.ParallelConfig.set_dcp_defaults].
363
+ """
364
+
365
+ dcp_q_replicate: bool | None = None
366
+ """Replicate the MLA query projection within each DCP group so decode can skip the
367
+ query all-gather.
368
+
369
+ With DCP the KV cache is sharded across the group, so the standard MLA decode path
370
+ all-gathers the query every step. Replicating the (small) query projection at load
371
+ time lets each rank materialize the full group-local head set and skip that
372
+ collective, at the cost of computing the projection redundantly on every rank
373
+ in the group.
374
"""
375
376
cp_kv_cache_interleave_size: int = 1
@@ -538,13 +552,23 @@ class ParallelConfig:
552
f"{sorted({1, pcp, tp * pcp})}."
553
)
554
541
- if self.dcp_comm_backend == "a2a" and self.decode_context_parallel_size <= 1:
542
- raise ValueError(
543
- "dcp_comm_backend='a2a' requires decode_context_parallel_size > 1."
544
- )
545
-
555
return self
556
557
+ def set_dcp_defaults(
558
+ self,
559
+ comm_backend: DCPCommBackend = "ag_rs",
560
+ q_replicate: bool = False,
561
+ ) -> None:
562
+ """Fill in the DCP options the user left unset.
563
+
564
+ Models can set their preferred DCP settings by calling this from their
565
+ `verify_and_update_config` hook.
566
+ """
567
+ if self.dcp_comm_backend is None:
568
+ self.dcp_comm_backend = comm_backend
569
+ if self.dcp_q_replicate is None:
570
+ self.dcp_q_replicate = q_replicate
571
+
572
@property
573
def world_size_across_dp(self) -> int:
574
"""Process world size across TP, PCP, PP, and DP."""
vllm/config/vllm.py
+4
@@ -1126,6 +1126,10 @@ class VllmConfig:
1126
1127
self.try_verify_and_update_config()
1128
1129
+ # Models may have supplied their own DCP defaults above; anything still
1130
+ # unset falls back to the stock ones.
1131
+ self.parallel_config.set_dcp_defaults()
1132
+
1133
if self.model_config is not None:
1134
self.model_config.verify_with_parallel_config(self.parallel_config)
1135
self.model_config.verify_dual_chunk_attention_config(self.load_config)
vllm/engine/arg_utils.py
+7
-1
@@ -481,7 +481,8 @@ class EngineArgs:
481
tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
482
prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
483
decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
484
- dcp_comm_backend: DCPCommBackend = ParallelConfig.dcp_comm_backend
484
+ dcp_comm_backend: DCPCommBackend | None = ParallelConfig.dcp_comm_backend
485
+ dcp_q_replicate: bool | None = ParallelConfig.dcp_q_replicate
486
dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
487
cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
488
data_parallel_size: int = ParallelConfig.data_parallel_size
@@ -1070,6 +1071,10 @@ class EngineArgs:
1071
"--dcp-comm-backend",
1072
**parallel_kwargs["dcp_comm_backend"],
1073
)
1074
+ parallel_group.add_argument(
1075
+ "--dcp-q-replicate",
1076
+ **parallel_kwargs["dcp_q_replicate"],
1077
+ )
1078
parallel_group.add_argument(
1079
"--dcp-kv-cache-interleave-size",
1080
**parallel_kwargs["dcp_kv_cache_interleave_size"],
@@ -2295,6 +2300,7 @@ class EngineArgs:
2300
worker_extension_cls=self.worker_extension_cls,
2301
decode_context_parallel_size=self.decode_context_parallel_size,
2302
dcp_comm_backend=self.dcp_comm_backend,
2303
+ dcp_q_replicate=self.dcp_q_replicate,
2304
dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,
2305
cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
2306
_api_process_count=self._api_process_count,
vllm/model_executor/models/config.py
+11
@@ -40,6 +40,16 @@ class DeepseekV32ForCausalLM(VerifyAndUpdateConfig):
40
logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
41
42
43
+class GlmMoeDsaForCausalLM(VerifyAndUpdateConfig):
44
+ @staticmethod
45
+ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
46
+ # For Glm-Moe-DSA, qrep + a2a is better than the default all-gather + ag-rs
47
+ # in most cases.
48
+ vllm_config.parallel_config.set_dcp_defaults(
49
+ comm_backend="a2a", q_replicate=True
50
+ )
51
+
52
+
53
class Ernie4_5_VLMoeForConditionalGenerationConfig(VerifyAndUpdateConfig):
54
@staticmethod
55
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
@@ -913,6 +923,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
923
"Gemma4ForCausalLM": Gemma4Config,
924
"Gemma4ForConditionalGeneration": Gemma4Config,
925
"Gemma4UnifiedForConditionalGeneration": Gemma4Config,
926
+ "GlmMoeDsaForCausalLM": GlmMoeDsaForCausalLM,
927
"GptOssForCausalLM": GptOssForCausalLMConfig,
928
"LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig,
929
"GteModel": SnowflakeGteNewModelConfig,
vllm/model_executor/models/deepseek_v2.py
+7
-1
@@ -1041,8 +1041,14 @@ class DeepseekV2MLAAttention(nn.Module):
1041
prefix=f"{prefix}.kv_a_proj_with_mqa",
1042
)
1043
1044
- qrep_enabled = (
1044
+ # The env var predates the config field and still wins if set explicitly.
1045
+ qrep_requested = (
1046
envs.VLLM_DCP_Q_REPLICATE
1047
+ if envs.is_set("VLLM_DCP_Q_REPLICATE")
1048
+ else bool(vllm_config.parallel_config.dcp_q_replicate)
1049
+ )
1050
+ qrep_enabled = (
1051
+ qrep_requested
1052
and vllm_config.parallel_config.decode_context_parallel_size > 1
1053
and vllm_config.parallel_config.prefill_context_parallel_size <= 1
1054
)