diff --git a/tests/distributed/test_dcp_a2a.py b/tests/distributed/test_dcp_a2a.py index 73b6e4a18b..596fec5d77 100644 --- a/tests/distributed/test_dcp_a2a.py +++ b/tests/distributed/test_dcp_a2a.py @@ -105,20 +105,11 @@ class TestDCPCommBackendConfig: """Test --dcp-comm-backend config validation.""" def test_default_is_ag_rs(self): - """Default comm backend is ag_rs.""" + """Comm backend resolves to ag_rs unless the model asks otherwise.""" config = ParallelConfig() + config.set_dcp_defaults() assert config.dcp_comm_backend == "ag_rs" - def test_a2a_requires_dcp_greater_than_1(self): - """A2A backend requires decode_context_parallel_size > 1.""" - with pytest.raises( - ValueError, match="requires decode_context_parallel_size > 1" - ): - ParallelConfig( - dcp_comm_backend="a2a", - decode_context_parallel_size=1, - ) - def test_a2a_with_dcp_valid(self): """A2A backend is valid when DCP > 1.""" config = ParallelConfig( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index d0118c6503..15c0116e4d 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -351,12 +351,26 @@ class ParallelConfig: and will be deprecated when PCP is fully supported. """ - dcp_comm_backend: DCPCommBackend = "ag_rs" + dcp_comm_backend: DCPCommBackend | None = None """Communication backend for Decode Context Parallel (DCP). - - "ag_rs": AllGather + ReduceScatter (default, existing behavior) + - "ag_rs": AllGather + ReduceScatter (existing behavior) - "a2a": All-to-All exchange of partial outputs + LSE, then combine with Triton kernel. Reduces NCCL calls from 3 to 2 per layer for MLA models. + + `None` selects the model default, which is "ag_rs" unless the model + overrides it via [`set_dcp_defaults`][vllm.config.ParallelConfig.set_dcp_defaults]. + """ + + dcp_q_replicate: bool | None = None + """Replicate the MLA query projection within each DCP group so decode can skip the + query all-gather. + + With DCP the KV cache is sharded across the group, so the standard MLA decode path + all-gathers the query every step. Replicating the (small) query projection at load + time lets each rank materialize the full group-local head set and skip that + collective, at the cost of computing the projection redundantly on every rank + in the group. """ cp_kv_cache_interleave_size: int = 1 @@ -538,13 +552,23 @@ class ParallelConfig: f"{sorted({1, pcp, tp * pcp})}." ) - if self.dcp_comm_backend == "a2a" and self.decode_context_parallel_size <= 1: - raise ValueError( - "dcp_comm_backend='a2a' requires decode_context_parallel_size > 1." - ) - return self + def set_dcp_defaults( + self, + comm_backend: DCPCommBackend = "ag_rs", + q_replicate: bool = False, + ) -> None: + """Fill in the DCP options the user left unset. + + Models can set their preferred DCP settings by calling this from their + `verify_and_update_config` hook. + """ + if self.dcp_comm_backend is None: + self.dcp_comm_backend = comm_backend + if self.dcp_q_replicate is None: + self.dcp_q_replicate = q_replicate + @property def world_size_across_dp(self) -> int: """Process world size across TP, PCP, PP, and DP.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 78d7d083b8..a3e0276176 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1126,6 +1126,10 @@ class VllmConfig: self.try_verify_and_update_config() + # Models may have supplied their own DCP defaults above; anything still + # unset falls back to the stock ones. + self.parallel_config.set_dcp_defaults() + if self.model_config is not None: self.model_config.verify_with_parallel_config(self.parallel_config) self.model_config.verify_dual_chunk_attention_config(self.load_config) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 293e67be5d..6c574d3238 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -481,7 +481,8 @@ class EngineArgs: tensor_parallel_size: int = ParallelConfig.tensor_parallel_size prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size - dcp_comm_backend: DCPCommBackend = ParallelConfig.dcp_comm_backend + dcp_comm_backend: DCPCommBackend | None = ParallelConfig.dcp_comm_backend + dcp_q_replicate: bool | None = ParallelConfig.dcp_q_replicate dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size data_parallel_size: int = ParallelConfig.data_parallel_size @@ -1070,6 +1071,10 @@ class EngineArgs: "--dcp-comm-backend", **parallel_kwargs["dcp_comm_backend"], ) + parallel_group.add_argument( + "--dcp-q-replicate", + **parallel_kwargs["dcp_q_replicate"], + ) parallel_group.add_argument( "--dcp-kv-cache-interleave-size", **parallel_kwargs["dcp_kv_cache_interleave_size"], @@ -2295,6 +2300,7 @@ class EngineArgs: worker_extension_cls=self.worker_extension_cls, decode_context_parallel_size=self.decode_context_parallel_size, dcp_comm_backend=self.dcp_comm_backend, + dcp_q_replicate=self.dcp_q_replicate, dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size, cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, _api_process_count=self._api_process_count, diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d294943318..12c9f5630e 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -40,6 +40,16 @@ class DeepseekV32ForCausalLM(VerifyAndUpdateConfig): logger.info("Using bfloat16 kv-cache for DeepSeekV3.2") +class GlmMoeDsaForCausalLM(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + # For Glm-Moe-DSA, qrep + a2a is better than the default all-gather + ag-rs + # in most cases. + vllm_config.parallel_config.set_dcp_defaults( + comm_backend="a2a", q_replicate=True + ) + + class Ernie4_5_VLMoeForConditionalGenerationConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: @@ -913,6 +923,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "Gemma4ForCausalLM": Gemma4Config, "Gemma4ForConditionalGeneration": Gemma4Config, "Gemma4UnifiedForConditionalGeneration": Gemma4Config, + "GlmMoeDsaForCausalLM": GlmMoeDsaForCausalLM, "GptOssForCausalLM": GptOssForCausalLMConfig, "LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig, "GteModel": SnowflakeGteNewModelConfig, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index d8bc3fb64b..93ce2c7f3d 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1041,8 +1041,14 @@ class DeepseekV2MLAAttention(nn.Module): prefix=f"{prefix}.kv_a_proj_with_mqa", ) - qrep_enabled = ( + # The env var predates the config field and still wins if set explicitly. + qrep_requested = ( envs.VLLM_DCP_Q_REPLICATE + if envs.is_set("VLLM_DCP_Q_REPLICATE") + else bool(vllm_config.parallel_config.dcp_q_replicate) + ) + qrep_enabled = ( + qrep_requested and vllm_config.parallel_config.decode_context_parallel_size > 1 and vllm_config.parallel_config.prefill_context_parallel_size <= 1 )